[OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head - #3131

…3134)

> [!NOTE]
> **Slice [1/3] of #2746's split** (2026-10-30), so the model-spec work
is reviewed and merged in pieces under the 500-line budget. Merge order:
#2741 [1/3] (#3131, merged) → **this PR** → #2746 [2/3] (#3135) → #2746
[3/3] (#2746).
>
> **Base:** `main`, since #2741 [1/3] (#3131) merged on 2026-10-30. The
diff below is this PR's alone.

<details><summary>Series layout (the part around this PR)</summary>

- `main` (#2741 [1/3], #3131, merged)
- **#2746 [1/3] the expert-parallel sections of the model specs (this
PR)**
    - #2746 [2/3] (#3135) the spec applier
  - the expert-parallel tree: #3136 #2741 [2/3] → #2780
- #2746 [3/3] (#2746): the decision and the load apply the spec. It sits
on a helper branch that joins #2780 and #2746 [2/3].

</details>

### What does this PR do?

Type of change: New feature

Before it can shard a model's experts, the expert-parallel path needs
three facts about the model:
- which modules are the routed experts, and how they shard;
- which transformers release first ships that information;
- for a VLM, which nested config the plan must be attached to.

All three are facts about the architecture, so they belong with the
model's other facts in `modelopt.torch.models`.

- **`EPSpec`**, a new `ModelSpec` section (`ModelSpec.ep_spec`). Its
fields:
  - `base_model_ep_plan`, the plan to attach;
  - `plan_config_attr`, the nested config it goes on;
- `requires_transformers`, the release this model's expert parallelism
needs;
  - `prepare`, for a checkpoint that ships its own modeling code;
  - `prepared_ep_plan`, the plan `prepare` attaches;
  - `unsupported`, why a model cannot be expert-parallel yet.
- **Most models become data in the spec package they already had:**
- `deepseek_v3` carries the plan stock transformers ships only from
5.13;
- `qwen3_moe` and `minimax_m3_vl` carry only the transformers release
their own plan needs (a new `minimax_m3_vl` package);
- `nemotron_h` carries its plan. It needs no `prepare`: its gate emits
logits and the block runs its own top-k, so the experts receive global
expert ids, which the all-to-all experts forward routes on directly.
- **One thing stays code:** `EPSpec.prepare`, the escape hatch for
checkpoints that ship their own modeling code (the Kimi adapters). A
spec whose plan only `prepare` attaches states that plan too
(`prepared_ep_plan`). It can then be read without running `prepare`,
which patches the model's code.
- **Nemotron-H-Omni** gets its own spec package.
- Its experts live in the language sub-model, and transformers
aggregates a sub-model's plan up to the root with its module prefix. A
plan on the VLM root would match nothing, silently.
- The plan goes on the language-model config. The checkpoint's own code
names it `llm_config` and transformers' 5.18 port `text_config`. The
spec lists both, `text_config` first, and the first one a config has
takes the plan.
- transformers ships the Omni model itself from 5.18, which its spec
records. Its root type registers only the expert-parallel plan, so
`is_moe` by the root type is False for its `NemotronHMoE` blocks.
- No export depends on the missing MoE and export sections: on
transformers 5 its experts are one fused container, exported by its own
handler (checked locally, see Testing).
- **`EPSpec.unsupported`** says why a model whose own plan would route
it into the expert-parallel path cannot take it yet:
- GPT-OSS and Llama-4 store their experts transposed, which the
fused-experts export cannot write. GPT-OSS also interleaves gate and up.
- Llama-4's MoE block also routes tokens itself rather than through the
experts interface. A new `llama4_text` package shares Llama-4's reason.

Nothing reads the section yet. #2746 [2/3] (#3135) applies it to a
config, and #2746 [3/3] (#2746) has the expert-parallel decision and
load use it.

### Usage

```python
from modelopt.torch.models import get_spec

get_spec("deepseek_v3").ep_spec.base_model_ep_plan  # the plan transformers ships only from 5.13
get_spec("nemotron_h_omni").ep_spec.plan_config_attr  # ("text_config", "llm_config")
get_spec("gpt_oss").ep_spec.unsupported  # why GPT-OSS is not expert-parallel yet
```

### Testing

- New `tests/unit/torch/models/test_ep_specs.py`, the registered specs:
- every registered plan is well formed, and every EP section declares
something;
- Nemotron-H-Omni shares Nemotron-H's plan and puts it on its
language-model config.
- `test_model_specs.py` and `test_specs_vs_transformers.py`, which check
every registered spec, pass with the new packages.
- Checked locally, not as a unit test: an FP8, INT4-AWQ or W4A8-AWQ
export of a tiny NemotronH writes byte-identical tensors with and
without nemotron_h's MoE and export sections. Environments: transformers
5.18 and 5.5.4 on torch 2.14, and 5.18 on torch 2.8.
- Declined review suggestion (on #2746, 2026-10-28): checking the
version claims of specs without a `moe_spec` is left for the PR that
next touches `test_specs_vs_transformers.py`, as the review suggests;
the gap also covers main's own such specs. The three claims added here
hold on transformers 5.4, 5.5.4, 5.12, 5.18 and 5.19.
- 2026-10-30, at this head: pre-commit, mypy included, passes against
main. CPU tests (`tests/unit/torch/models` and `tests/unit/torch/utils`,
torch 2.14):

| transformers | Passed | Skipped | Failed |
|---|---|---|---|
| 5.18 | 751 | 0 | 0 |
| 5.5.4 | 734 | 17 | 0 |
- Rebased on 2026-10-30 (fe2e966) onto main 99c3834 through
#3131's rebase; no conflicts, this PR's own diff unchanged. Pre-commit
passes against main.
- Review round 2, 2026-10-30 (bc55734, 3b43cc6): the
Nemotron-H-Omni spec holds its own copy of Nemotron-H's plan, and the
reuse test checks the two specs do not share one dict (it fails without
the copy); the text-only Llama-4 spec's docstring says why it registers
no MoE section. Pre-commit passes against main;
`tests/unit/torch/models` and `test_moe_a2a.py` pass on transformers
5.18 (247) and 5.5.4 (241, 6 skipped).
- Rebased on 2026-10-30 (d01b7b2) onto main 62a5c48, after #3131
merged: no conflicts, this PR's own diff unchanged (range-diff).
Pre-commit passes against main; `tests/unit/torch/models` and
`test_moe_a2a.py` pass on transformers 5.18 (286) and 5.5.4 (268, 7
skipped).

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ Additive: a new spec section
and three new spec packages, which nothing reads until #2746 [3/3].
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A. The feature's entry lands with #2754.
- Did you get Claude approval on this PR?: ❌ Not yet. Requested; #2746
as a whole was approved on 2026-10-28.

### Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models. Split out of
#2746 (OMNIML-6076), whose review history stays on #2746.





<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

<!-- coderabbit-ai-generated:v1 -->

## Summary by CodeRabbit

- **New Features**
- Added expert-parallel PTQ support information for DeepSeek-V3,
Nemotron-H, Nemotron-H Omni, and Qwen3-MoE, including model-specific
plan and Transformers version requirements.
- Added model registrations for Llama-4 Text, MiniMax-M3-VL, and
Nemotron-H Omni, with applicable Transformers requirements.
- Clarified when expert-parallel support is unavailable for GPT-OSS and
Llama-4.
- Added an option to load supported vision-language and encoder-decoder
models without requiring a language-model head.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…2a instead of forcing eager (#3074)

> [!NOTE]
> **[2/2]**: #3122 ([1/2]) is this PR's expert-parallel correctness fix,
split out on 2026-10-29 at a reviewer's request so that it lands on its
own, with a Bug Fixes changelog entry. This PR keeps the two
backward-breaking changes, calibration on `grouped_mm` by default and
the new `calibrate()` error, and sat on #3122 until #3122 merged on
2026-10-30. Before that, it was [3/3] of the OMNIML-6073 split, after
#3088 (FSDP2's weight-access window) and #2745 (`grouped_mm`
quantization), both merged. #3067 (Nemotron-H) sits on this PR.
>
> **Base:** `main`. #3122 merged on 2026-10-30, so this PR sits on main,
its own diff unchanged.
>
> **In the tree:** #2741's helper base merges #3067, which carries this
PR and #2745.

<details><summary>Series layout</summary>

- `main`
  - #3088 FSDP2 weight-access window (merged)
  - #2745 fused experts under grouped_mm (merged)
- #3122 expert-parallel experts calibrated and folded on each rank's own
experts (merged)
  - **#3074 calibration on grouped_mm (this PR)**
    - #3067 Nemotron-H's experts under grouped_mm
  - #2739 meta-DCP checkpoint load (merged)
  - #2736 distribution plan, run group, collectives (merged)
    - #2740 calibration on per-rank data
    - #2750 keep-fused fold
    - #2751 distributed-writer pieces
      - #2752 no-gather export (on #2751, #2750, #2779 and #2740)
        - #2753 expert-parallel export
    - #2737 plan decision
      - #2779 load routing
        - #2741 expert-parallel loader (on #2779, #3067 and #2740)
          - #2742 dense sharding
            - #2743 block-FP8 dequantization
              - #2780 expert-parallel load routing
                - #2746 expert parallelism in the model specs
                  - #2747 on-the-fly decompression
                    - #2748 Kimi-K2.5 and Kimi-VL adapters
                      - #2749 Kimi-K3 adapter
- #2754 expert-parallel PTQ from hf_ptq.py (on #2746 and #2753)

</details>

### What does this PR do?

Type of change: New feature, with a backward-incompatible change to MoE
calibration (see the changelog entry)

#2745 quantizes the `grouped_mm` experts forward, transformers' default,
while calibration still runs every model on eager. This PR moves
calibration onto `grouped_mm` and the `torch_a2a` expert-parallel
forward, and makes everything that reads the experts' inputs or weights
during calibration work there. Each gap below otherwise failed silently,
by collecting no amax or the wrong one rather than by raising.

- **Backend normalization.** `normalize_experts_impl_on_the_fly`
replaces `force_eager_experts_impl_on_the_fly`. `eager`, `grouped_mm`
and `torch_a2a` are kept; a backend nothing intercepts (`batched_mm`,
the opaque `sonicmoe` kernel, hub kernels) is moved to `grouped_mm`
before calibration, or to `eager` on a transformers without
`_grouped_linear`. A model whose experts quantize only on the eager
forward keeps main's forced eager: native FP8 experts, a class with its
own wrapper (GPT-OSS's quantizes `down_proj`'s input by intercepting the
eager forward's matmuls), and Nemotron-H's non-gated experts until
#3067. Such a model is refused on `torch_a2a`, before any config
changes, instead of calibrating on raw FP8 bytes or not at all.
- **Per-expert captures.** `grouped_mm` and `torch_a2a` feed a
projection's input quantizer one expert-sorted batch. The weight-input
captures (local Hessian, imatrix, GPTQ) split it back per expert by the
cumulative token offsets the grouped interception records, following
whichever path ran. The input quantizers calibrate on the rows the
grouped GEMM computes only: under transformers' own expert parallelism,
the rows routed to other ranks' experts are sorted past the offsets, and
before transformers 5.17 the down projection's input holds them
uninitialized.
- **Expert parallelism.** #3122 makes everything that reads the expert
weights outside the forward index each rank's own experts: max
calibration's weight pass, `fold_weight`, the weight-input captures and
GPTQ. Here the grouped captures split the expert-sorted batch per local
expert. The grouped weight quantization needs no indexing: transformers
swaps expert-parallel DTensors for their local shards around the experts
forward.
- **Fail loudly.** `calibrate()` raises, naming the modules, when a
fused-experts forward that had to reach a quantizer ran but neither hook
fired: a hand-written `bmm` forward on a class without
`@use_experts_implementation`, which no config change redirects. Each
`calibrate()` judges its own forwards and raises with the model back in
its training mode. Under the expert parallelism of transformers 5.16 and
later a module's verdict is reduced over its experts' expert-parallel
group, so every rank of the group raises, not just the ranks whose
experts saw rows. Not reported: experts no calibration token reached
(Ernie-4.5-VL's vision experts under text-only calibration, which fall
back to weight-derived scales as on main), an eager forward whose rows
were all routed to other ranks' experts, and weight-only experts under
max or MSE, whose weight quantizers calibrate on the weights alone.

**Behavior change.** main forced MoE calibration onto eager;
transformers resolves an unrequested backend to `grouped_mm`, so
fused-experts models now calibrate there. A single-GPU Qwen3-30B-A3B
export, NVFP4 or FP8, matches main in every weight and weight scale, but
4,647 activation scales move, by up to 10.2%. main's own repeat is
byte-identical, and loading with `experts_implementation="eager"`
reproduces main's checkpoint byte for byte (all 74,835 tensors for NVFP4
and 56,211 for FP8), so the backend alone moves them. The changelog says
so.

| Moved activation scales, Qwen3-30B-A3B NVFP4, 1 GPU | Tensors |
|---|---|
| Experts' `down_proj` input (33 layers, 128 experts each) | 4,224 |
| Experts' shared `gate_proj` and `up_proj` input (1 layer) | 256 |
| Attention input and KV-cache | 167 |
| Total, in 45 of 48 layers (none in layers 0, 1 and 5) | 4,647 |

The largest moves are 10.2% on the experts (layer 25's `down_proj`
input) and 7.8% on attention.

*Why they move.* Both forwards feed an input quantizer the same rows: a
single GPU has no padding rows (transformers' sentinel rows exist only
under its own expert parallelism, where calibration skips them;
transformers zeroes them only from 5.17), the routing weights apply
after `down_proj` on both, and neither casts. They round the block's
output differently. Eager adds a token's 8 expert outputs into a bf16
buffer with `index_add_`, rounding after each add, while
`grouped_mm_experts_forward` sums them in fp32 and rounds once, and its
GEMMs run other kernels. From the first MoE block's output on, the
residual stream differs by rounding, and where a router's 8th and 9th
scores nearly tie, the routing flips. A flipped token feeds `down_proj`
another expert's activations, and `down_proj`'s input scale is a max
over routed (token, expert) rows, so it moves by more than rounding
would. The moved scales fit that: the experts' shared
`gate_proj`/`up_proj` input, whose rows are every token however it is
routed, moved in one layer only, against 33 for `down_proj`.

*Why `grouped_mm` off expert parallelism too,* rather than eager there
and `grouped_mm` on `torch_a2a` only: it is the forward transformers
runs the model on when nothing else is asked for, its fp32 sum is the
more precise of the two combines, and expert-parallel calibration runs
the grouped path anyway (`torch_a2a` runs the local experts through
`_grouped_linear`), so a model calibrates on one forward whether or not
it is sharded. `experts_implementation="eager"` keeps earlier releases'
checkpoints.

### Usage

```python
# To calibrate exactly as earlier releases did:
model = AutoModelForCausalLM.from_pretrained(path, experts_implementation="eager")
```

```bash
# The same from hf_ptq.py:
python hf_ptq.py --pyt_ckpt_path <model> --qformat nvfp4 --experts_implementation eager  # or grouped_mm
```

### Testing

- `tests/unit/torch/quantization/plugins/test_fused_experts.py`, against
synthetic and real transformers experts modules:
- backend normalization, including the fallback to eager on a
transformers without `_grouped_linear`;
- GPTQ's Hessians and local Hessian's captures under `grouped_mm` and
`torch_a2a` match eager's per expert, and a capture follows whichever
path ran;
- the interception check raises when the forward ran past both hooks,
also after an earlier calibration reached them, and not for unreached
experts, an eager forward routed only to other ranks' experts, or
weight-only experts beside a calibrated dense input quantizer under max
(each expert's weight amax equals its weight's per-channel abs max).
Under local Hessian and GPTQ, which read the experts' inputs, those
weight-only experts still raise;
- GPT-OSS, a real tiny model, calibrating `down_proj`'s input on the
eager forward.
- Expert parallelism, on CPU under gloo:
- the tests that experts sharded the way transformers shards them
calibrate, fold, capture and build GPTQ helpers on each rank's own
experts moved to #3122 with that fix;
- every rank of the experts' expert-parallel group raises when only one
rank's experts ran past both hooks;
- with fake quantizers inserted, experts split over two ranks train
through the `torch_a2a` forward like the whole quantized module: the
fake-quantized case of `test_moe_a2a.py`'s whole-module test;
- the transport's hop slicing under a graph: `test_moe_a2a.py`'s
chunked-hops test also runs forward and backward through hops sliced
below the limit, and checks the output and input gradients bit for bit
against the unsliced run.
- The transport's (#2738, merged) docs say again that ModelOpt quantizes
its local experts, now that the `_grouped_linear` override does. Its
hop-slicing comments say what the slicing guarantees, every hop within
`all_to_all_single`'s int32 element offsets, rather than that it
prevents a hang: with the slicing disabled, hops past 2**31 elements
completed exactly on NCCL 2.29, 2.30 and 2.31, including at 32 ranks
over the network on 2.29 and 2.30.
- Against main, on tiny real models: GPT-OSS calibrates on eager with
`down_proj`'s input amax equal to main's; Ernie-4.5-VL-MoE quantizes in
FP8 and NVFP4 under text-only calibration, as on main.
- Byte-for-byte checkpoint comparison against main on one GPU, measured
on earlier revisions of #2745: Qwen3-30B-A3B, eager and `grouped_mm`,
NVFP4 and FP8; Qwen2.5-0.5B (dense) is byte-identical. On GPU,
Nemotron-3-Nano-30B-A3B's NVFP4 calibration with its experts pinned to
eager and #3067's with `experts_implementation="eager"` give
byte-identical exports, all 18,012 tensors.
- The GPU tests of QAT under expert parallelism are in #2742, the first
PR with an FSDP2-sharded dense part they also run.
- Review fixes, 2026-10-27, carried from #2745: the FSDP2 window's
ownership test (unnamed meshes, every shard state); an empty experts
forward is no work the interception check expects; the non-gated experts
are pinned to eager whichever plugin runs first.
- Review fixes, 2026-10-27, last round, each with a test that fails
without it:
- `grouped_mm` is kept, and redirected to, only where `_grouped_linear`
can be intercepted; one capability constant gates the interception, the
supported set and the redirect target;
- an eager experts forward counts as work for the interception check
only when it routes a row to this rank's experts, so a rank whose
experts calibration routed nothing to no longer raises alone;
- an enabled expert weight quantizer needs the interception only when
the calibration captures the experts' inputs to calibrate their weights
(GPTQ, local Hessian, imatrix), so weight-only experts under max or MSE
no longer raise; the message names what falls back.
- Split from #2745 on 2026-10-27, night: #2745's and this PR's commits
together give #2745's previous tree exactly, plus this round's fixes.
Pre-commit, mypy included, and this PR's unit tests pass at this head on
transformers 5.18 and 5.5.4 and on torch 2.8, and the series' PRs merged
together add no unit-test failure over main f03651e
(tests/unit/torch utils, export, quantization and models). This
revision's GPU tests run in CI only.
- Review suggestions, 2026-10-28: the weight-input captures (local
Hessian, imatrix, GPTQ) read the path that ran, not the one the config
named when they were registered, so they follow whichever forward runs
and the refusal is gone, pinned by a test that fails on the code before;
the interception check's forward count, and the host sync its routing
test costs, stop after the first hook call of each calibration; the
check's docstring says it is rank-local. Rebased over #2745's fixes,
taking back the `torch_a2a` test fixtures #2745 no longer carries.
- Rebased over #2745's second-round fixes: the grouped interceptor
records the offsets and clears the eager expert index, and #2745's test
that captures skip the grouped batch gives way to this PR's per-expert
capture tests.
- Review follow-up, 2026-10-28, third round: drops the note #2745 now
adds, that the grouped captures skip their batch, since here they split
it per expert. A new test pins that a calibration forward stops
counting, and skips its host-syncing routing test, once a hook has
fired; it fails with that gate removed.
- Review suggestions, 2026-10-28, fourth round: the experts forward
clears `_grouped_offs` when it returns, so a capture reached outside a
forward skips instead of splitting its input by a finished forward's
offsets; the interception check's docstring says that a rank whose
experts saw no rows waits in its next collective until the launcher
stops it or the collective times out, so the error is in another rank's
log. Rebased over #2745's cleanup.
- Rebased on 2026-10-28, evening, onto #2745's head `e7c449fcf87`, after
#2745 moved onto main `7092db6ac5b` and took its last review rounds.
Conflicts only in the changelog (both sides' entries kept),
`_QuantFusedExperts`'s docstring (it keeps #2745's note that the
`grouped_mm` fallback does not call `F.linear`) and the fused-experts
tests (#2745's transposed-container tests kept). The FP8 experts
wrapper's comment from #2745 now names
`normalize_experts_impl_on_the_fly`, this PR's name for the function
that keeps it on eager. pre-commit, mypy included, passes at this head
against main.
- Review fix, 2026-10-28, evening: `calibrate()`'s interception check
reads the HF plugin from `sys.modules` instead of importing it, so a
plugin that failed to load (which `import_plugin` only warns about) no
longer fails every calibration with its error.
`test_a_plugin_that_failed_to_load_leaves_calibration_alone` breaks the
plugin's import and fails without the change. pre-commit, mypy included,
passes at this head.
- Split, 2026-10-28, evening, at the bot review's request: the FSDP2
window's ownership filter and its ownership test moved to #3088, with a
test that the window leaves an ignored expert-parallel DTensor sharded.
This PR sits on a helper base merging #2745 and #3088, and its source
diff is 334 lines. At this head `test_fused_experts.py`,
`test_moe_a2a.py` and `test_utils.py` pass on CPU (200 passed, 6
skipped), and pre-commit, mypy included, passes against main.
- Review suggestions, 2026-10-28, evening: `_local_experts` refuses a
DTensor whose local shard does not hold the module's experts (one
replicated over the ranks, say) instead of indexing the global experts,
and the native FP8 wrapper's `unpack_weight` refuses expert weights
sharded over the experts instead of filling a buffer for every expert.
One two-rank test covers both and fails without either. The suggested
all-reduce in the interception check is declined (a collective in every
distributed `calibrate()`). Focused tests pass (201 passed, 6 skipped),
and pre-commit passes against main.
- Review suggestions, 2026-10-28, late evening: the per-expert split of
a grouped batch requires exactly one exclusive end per expert, within
the rows, as transformers' `grouped_mm` and `a2a_moe` build the offsets,
and raises on any other form (a leading zero would credit each expert
with the previous one's rows). A test feeds three such forms and fails
without the check. The `_grouped_offs` comment names all three captures
and the attribute's role as the which-forward-ran signal. Rebased onto
the helper base re-merged with #3088's update; focused tests pass (204
passed, 6 skipped), and pre-commit passes against main.
- Review suggestion, 2026-10-28, late evening: `_local_experts` checks
for a DTensor through the plugin's guarded `torch.distributed.tensor`
import, so a torch without distributed does not fail every fused-MoE
call. Rebased onto the helper base re-merged with #3088's update;
focused tests pass (204 passed, 6 skipped), and pre-commit passes
against main.
- 2026-10-28, night: `hf_ptq.py` gains `--experts_implementation`,
passed on every load path: `get_model` (the compressed-tensors branch
included), the low-memory load, and the FSDP2 load, through the config
its meta model is built from. The changelog's way back to earlier
releases' scales is now reachable from the script. A test loads a tiny
Qwen3-MoE with `eager` and fails without the plumbing. The
`tests/examples/hf_ptq` unit tests pass (110), and pre-commit passes
against main.
- Review fixes, 2026-10-28, night: on the FSDP2 path the flag is set on
the pre-fetched config's `_experts_implementation`, as
`build_meta_causal_lm` sets `attn_implementation`, and a test captures
the config the FSDP2 loader receives (it fails without the assignment);
`get_model` passes the flag only when given, on every branch.
`calibrate()`'s new error for a fused-experts forward that reaches
neither hook has its own changelog entry. Rebased onto #2745, itself
rebased onto main `bc0b6e03237` after #3088 merged; focused and
`tests/examples/hf_ptq` tests pass (316 passed, 6 skipped), and
pre-commit passes against main.
- Review suggestions and a rebase, 2026-10-28, late night: the backend
normalization warns once for each change it makes to a caller's choice,
but not for a `None` backend pinned to eager, which already dispatches
to the eager forward. Rebased onto #2745 on main `e9f432affbb` (#3091):
`DTensor` now comes from the plugin's plain `torch.distributed.tensor`
import, since the guarded import is gone with the torch 2.10 floor.
Focused and `tests/examples/hf_ptq` tests pass (325 passed, 6 skipped),
and pre-commit passes against main.
- Review suggestion and a rebase, 2026-10-28, late night: the
backend-redirect notice goes through `warn_rank_0`, so an
expert-parallel run prints one copy. Rebased onto #2745's comment and
warning-text update; `_is_fused_experts_module`'s new docstring names
`normalize_experts_impl_on_the_fly`, which keeps GPT-OSS on eager here.
`test_fused_experts.py` and `test_utils.py` pass (197), and pre-commit
passes against main.
- Bot review of `2e253a509d0`, 2026-10-28, late night: the FSDP2 test
runs a Qwen3.5-MoE config too and checks the text config its experts
dispatch on, which the one assignment reaches through transformers'
recursive `_experts_implementation` setter (5.5.4 and 5.18; a tiny
Qwen3.5-MoE built on meta dispatches its experts on `eager` with the
flag and on `grouped_mm` without it). The test fails when the value
reaches the root config only, and the comment at the assignment says the
setter carries it. Declined: resolving the routing indices' parameter
name from the forward's signature, since every transformers MoE block
passes them positionally. The `tests/examples/hf_ptq` unit tests pass on
5.18 and 5.5.4 (120), and pre-commit passes against main.
- Bot review of `042e4bf1e55` (approved), 2026-10-28, late night: the
interception error names the ways forward: a class no backend change
redirects needs its own quantized wrapper, and meanwhile the rest of the
model quantizes with these modules' quantizers disabled, which the check
honours. Declined: caching the grouped offsets' host copy across the two
projections' captures (one small sync per experts module per batch under
GPTQ, local Hessian and imatrix, whose per-expert matmuls dominate).
pre-commit passes against main.
- Rebased on 2026-10-28, late night, onto main `1ca8bd1fc4d` after #2745
merged (squashed with the tree of its last head, so this PR's tree is
unchanged); the base is now `main`, so the GPU tests run. pre-commit
passes against main.
- CI fix, 2026-10-28, late night: the first run on main failed two cases
of `tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py` in the
trtllm `hf_ptq` lane (2 failed, 267 passed): its stand-in for hf_ptq's
parsed arguments lacked the new `experts_implementation`, which
`load_model` now reads. The stand-in gains it; the file passes on
transformers 5.18 and 5.5.4.
- Rebased on 2026-10-28, late night, onto main `bdc45368d10` (#2736
merged), with its own diff unchanged. Its tests with #3067's, #2736's
`test_distributed.py` and the `hf_ptq` unit tests pass (376), and
pre-commit passes against main.
- CI, 2026-10-29: the first full run on main at `df2e3fe4a76`, GPU lanes
included, passes.
- Split, 2026-10-29, at a reviewer's request: the expert-parallel
indexing fix, its tests and a Bug Fixes changelog entry moved to #3122
([1/2], on main), and this PR now sits on it, its own diff otherwise
unchanged. The review's three smaller items are taken: the interception
check reduces its verdict over the experts' expert-parallel group (a
two-rank test fails with the rank-local check); a test checks that every
experts class transformers decorates with `use_experts_implementation`
takes the routing indices second (transformers 5.5.4 and 5.18); and
`--experts_implementation` offers `eager` and `grouped_mm` as argparse
choices. pre-commit, mypy included, passes against main; this PR's
tests, `test_moe_a2a.py` and the `hf_ptq` unit tests pass (315).
- Review suggestion (Claude), 2026-10-29: the interception check reads
the expert-parallel group from the experts' weight, a DTensor sharded
over the experts from transformers 5.16 on; from 5.6 to 5.15 the shards
are plain tensors that name no group, so the check stays rank-local
there, as before. The docstrings say so. Restacked onto #3122's two
follow-ups; pre-commit passes against main.
- Comment, 2026-10-29: at `calibrate()`'s gate for the interception
check, a sentence says that both conditions, the forward loop having run
and the HF plugin being loaded, must come out the same on every rank,
since under expert parallelism the check all-reduces over the experts'
group. pre-commit passes against main.
- Test fix, 2026-10-29:
`tests/gpu/torch/quantization/test_fsdp2_autoquant.py`'s 2-GPU KV
AutoQuantize test builds hf_ptq's arguments by hand and lacked the new
`experts_implementation`, which `load_model`'s FSDP2 branch reads. It
now passes `None`. The 1-GPU PR runners skip that test, and main's
nightly runs it. No other test or example builds hf_ptq's arguments by
hand on a path that reads the flag. pre-commit passes against main.
- Rebased on 2026-10-30 onto #3122 on main `99c3834f62c`, with its own
diff unchanged and no conflicts. pre-commit passes against main.
- Review (sychen52), 2026-10-30: with every weight quantizer collecting
statistics without quantizing (calibration), the grouped forward returns
the stored weight itself instead of building a weight-sized copy slice
by slice on each forward. A CPU test checks the identity and the
collected per-expert amax, and fails on the previous code. On CPU, one
call on a Qwen3-30B-A3B-sized `gate_up_proj` (128 × 1536 × 2048, bf16,
768 MiB) raised peak RSS by 755 MiB before and 0 MiB after; a GPU
`max_memory_allocated` comparison is queued. The changelog entry now
says the experts stay on `grouped_mm` for every forward after
`mtq.quantize` or `mto.restore` (QAT/QAD, AutoQuantize scoring,
fake-quant evaluation), with `experts_implementation="eager"` as the way
back. `test_fused_experts.py` and `test_moe_a2a.py` pass on transformers
5.18 (220) and 5.5.4 (214, 6 skipped); pre-commit passes against main.
- Review (kevalmorabia97), 2026-10-30: the three changes no test covered
now have one each, or are gone. The `_local_experts` call in
`_quantize_grouped_weight` is dropped as unreachable: the weight arrives
as this rank's local shard inside the experts forward. New tests pin
that a grouped forward reaching neither hook is counted whatever its
routing indices hold, and that a capture after a grouped forward does
not split by its offsets. The calibration pass-through calls only
quantizers that calibrate (Claude's suggestion), with a constant-amax
test case. Each test fails with its change reverted.
`test_fused_experts.py`, `test_moe_a2a.py` and `test_hf_ptq_args.py`
pass on transformers 5.18 (283) and 5.5.4 (277, 6 skipped); pre-commit
passes against main.
- GPU calibration peak memory, 2026-10-30 (sychen52's request):
`torch.cuda.max_memory_allocated` around the `mtq.quantize(...)` call
only, after `reset_peak_memory_stats`. Qwen3-30B-A3B on the default
`grouped_mm` backend, recipe `general/ptq/fp8_default-kv_fp8`,
`calib_size` 32, `batch_size` 1, one NVIDIA GB200 for both revisions
(torch 2.11, transformers 5.13.0, CUDA 13.1).

  | Revision | Calibration peak (bytes) | MiB |
  |---|---|---|
| Before the pass-through change (`3f3544435d8`) | 864,901,632 | 824.8 |
  | After (`27c5f77feac`) | 174,759,936 | 166.7 |
  | Drop | 690,141,696 | 658.2 |
- Rebased on 2026-10-30 onto main `55898ee4e3b` after #3122 merged,
dropping #3122's pre-squash commits; its own diff is unchanged, with no
conflicts. pre-commit passes against main.
- CI fix, 2026-10-30: the routing-order test read transformers' modeling
sources with the platform's default encoding, cp1252 on Windows, which
16 of them (`modeling_altclip.py`, say) cannot be decoded with; it now
reads them as UTF-8, and the windows unit job's only failure goes with
it. The `linux` job's one failure on the previous head, a NAS sorting
test this PR does not touch, passed on rerun at the same commit.
pre-commit passes against main.
- `eager` against main, byte for byte, 2026-10-30 (re-verifying at this
PR's head the claim that `experts_implementation="eager"` reproduces
main): PTQ of Qwen/Qwen3-30B-A3B with recipe
`general/ptq/nvfp4_default-kv_fp8`, `calib_size` 512, calibration
sequence length 2048, `batch_size` 1, cnn_dailymail, without expert
parallelism (one process), on one NVIDIA GB200 (torch 2.11, transformers
5.13.0, CUDA 13.1). This PR at `27c5f77feac` with
`--experts_implementation eager`, against main at `99c3834f62c`, both
compared at zero tolerance. Later revisions of this PR leave its own
diff unchanged (a rebase) or change tests only.

  | Scale kind | Tensors | Differing |
  |---|---|---|
  | `input_scale` | 18,624 | 0 |
  | `weight_scale` | 18,624 | 0 |
  | `weight_scale_2` | 18,624 | 0 |
  | `k_scale` | 48 | 0 |
  | `v_scale` | 48 | 0 |
- Review (juhi10071998), 2026-10-30: under transformers' own expert
parallelism the grouped forward sorts the rows routed to other ranks'
experts past `offs[-1]`, which the grouped GEMM neither reads nor
writes; before transformers 5.17 nothing zeroes them in the down
projection's input, so its input quantizer calibrated on uninitialized
rows. While an input quantizer collects statistics it now sees only the
rows before `offs[-1]` (a host sync, during calibration only). A test
hands the grouped projections a NaN or a large finite tail, as pre-5.17
expert parallelism leaves it, and checks the input amaxes match a run
without one; both cases fail without the change.
`test_fused_experts.py`, `test_moe_a2a.py`, `test_calib.py` and
`test_hf_ptq_args.py` pass on transformers 5.18 (307) and 5.5.4 (301, 6
skipped); pre-commit passes against main.
- Rebased on 2026-10-30 onto main `059588b89a5` (after #3131, #3124,
#2751, #3104, #3128, #3134 and #3149 merged), with its own diff
unchanged and no conflicts. pre-commit passes against main.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ❌ Fused-experts MoE calibration
now runs on `grouped_mm`, so some activation scales differ from earlier
releases; `experts_implementation="eager"` restores the earlier
checkpoint exactly.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ Backward Breaking Changes.
- Did you get Claude approval on this PR?: ✅ At `042e4bf1e55`; the
error-message change after it is from that review.

### Additional Information

Split out of #2745 (part 3 of 21 of expert-parallel PTQ for Hugging Face
MoE models), which merged on 2026-10-28.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
<!-- coderabbit-ai-generated:v1 -->

## Summary by CodeRabbit

* **Behavior Changes**
* Hugging Face MoE calibration now uses Transformers’ default
`grouped_mm` backend instead of forcing eager execution. Calibrated
activation scales may differ; set `experts_implementation="eager"` to
use the earlier calibration behavior.
* Eager-only expert wrappers continue to use eager execution. Using them
with `torch_a2a` raises an error.
* **Bug Fixes**
* Improved calibration support for grouped and expert-parallel
execution, including local expert weights and experts receiving no
routed inputs.
* Calibration now reports when fused experts run without interception or
required inputs are not captured.
* Improved FSDP2 parameter handling during weight access and writeback.
* **New Features**
* Added calibration support for newer Transformers fused-expert
implementations, including Qwen3-MoE and native FP8 experts.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…rom its spec (#3135)

> [!NOTE]
> **Slice [2/3] of #2746's split** (2026-10-30), so the model-spec work
is reviewed and merged in pieces under the 500-line budget. Merge order:
#2741 [1/3] (#3131, merged) → #2746 [1/3] (#3134) → **this PR** → #2746
[3/3] (#2746).
>
> **Base:** #2746 [1/3] (#3134): the spec sections this PR reads. It
also needs #2741 [1/3]'s `require_transformers`, on `main` since #3131
merged. The diff below is this PR's alone.

<details><summary>Series layout (the part around this PR)</summary>

- `main` (#2741 [1/3], #3131, merged)
  - #2746 [1/3] (#3134) the expert-parallel sections of the model specs
    - **#2746 [2/3] the spec applier (this PR)**
  - the expert-parallel tree: #3136 #2741 [2/3] → #2780
- #2746 [3/3] (#2746): the decision and the load apply the spec. It sits
on a helper branch that joins #2780 and this PR.

</details>

### What does this PR do?

Type of change: New feature

`modelopt.torch.utils.moe` gains the one applier of the model specs'
expert-parallel section (`EPSpec`, #2746 [1/3]), and the checks around
it.

- **`apply_distributed_support(model_type, hf_config, *, ckpt_path=None,
trust_remote_code=False)`** prepares a config for the expert-parallel
load from the model's spec, in this order:
1. It refuses a model its spec marks `unsupported`. The refusal names a
single process first, and FSDP2 (`--use_fsdp2` or
`MODELOPT_FSDP_SIZE=world`) only where the FSDP2 load supports the
checkpoint: an unquantized causal LM with no MTP head it cannot build.
GPT-OSS ships MXFP4, which the FSDP2 load refuses, Llama-4's checkpoint
is multimodal, and `--use_fsdp2` refuses Nemotron-H's MTP weights.
2. It checks the transformers floor: the expert-parallel load's own 5.7
(`_EP_MIN_TRANSFORMERS`), or the spec's later release. The load's floor
holds for a model with no spec too. Its refusal offers the same ways
around expert parallelism.
  3. It attaches the spec's plan to whichever config carries the MoE.
  4. It runs the spec's `prepare`, after the plan is attached.
- **An existing plan always wins** wherever the load reads it
(`_ep_config`): on the root, or on a nested config other than the one
the spec would take. Overwriting it would substitute our idea of the
layout for the model's.
- **A plan goes where the load reads it.**
- The applier refuses a spec whose `plan_config_attr` is neither
`text_config` nor `llm_config`, the two nested configs the
expert-parallel load reads a plan and the expert count from
(`_SUB_CONFIG_ATTRS`, which `_ep_config` and `_num_experts_from_config`
now read too). Attaching it elsewhere would be read as no plan at all
and plan the MoE as dense.
- `plan_config_attr` may list several spellings, the first one the
config has taking the plan. A config with none of them is refused,
naming them all.
- The spellings are checked for every spec, whether the spec or its
`prepare` supplies the plan, before anything is attached or prepared.
- **Read-only helpers:**
- `check_distributed_support` is the floor and refusal check alone, for
a config already prepared with its plan.
  - `_spec_refuses_ep` is its predicate.
- `_spec_ep_plan` is the plan a spec would supply: its own, or the one
it states `prepare` attaches, read without running `prepare`. A model
below its spec's transformers floor keeps its plan; a spec refusing
expert parallelism supplies none.

Nothing calls these yet. #2746 [3/3] (#2746) has the expert-parallel
decision and load call them.

### Usage

```python
from modelopt.torch.utils.moe import apply_distributed_support

hf_config = apply_distributed_support(hf_config.model_type, hf_config, ckpt_path=path)
# raises NotImplementedError for gpt_oss / llama4 / llama4_text, before any weight loads
```

### Testing

- `tests/unit/torch/models/test_ep_specs.py` gains the applier's tests:
- the applier's version floor, the load's floor for a model with no
spec, plan attachment, and `prepare` ordering;
- existing-plan precedence wherever the load reads the plan: the root,
or a nested config beside the one the spec would take;
- a spec attaching its plan where the load reads none is refused with
the config untouched, whether the spec or its `prepare` supplies the
plan;
- the Omni plan lands under either name (`llm_config` or `text_config`),
and on `text_config` for a config with both;
- GPT-OSS, Llama-4 and Llama-4-text are refused with their configs
untouched, naming a single process first, and a refusal comes before the
rest of the section;
  - the plan a spec would supply, read without running `prepare`.
- Mutation-checked: each of these fails a test:
  - attaching a plan where the load reads none;
- reading the Omni plan off `llm_config` alone, or trying `llm_config`
before `text_config`;
- checking a plan spelling only when the spec supplies the plan itself;
- checking for an existing plan on the spec's target rather than where
the load reads it;
  - hiding the plan of a spec below its transformers floor;
  - skipping the load's floor for a model with no spec;
  - dropping the EP-unsupported branch.
- Postponed review suggestion (on #2746, 2026-10-28): `EPSpec.prepare`
runs before the plan is known, so a load that ends up not
expert-parallel keeps its effects. No spec here sets `prepare`. It is to
be settled in #2748, which registers the first ones (Kimi-VL,
Kimi-K2.5).
- 2026-10-30, at this head: pre-commit, mypy included, passes against
main. CPU tests (`tests/unit/torch/models` and `tests/unit/torch/utils`,
torch 2.14):

| transformers | Passed | Skipped | Failed |
|---|---|---|---|
| 5.18 | 769 | 0 | 0 |
| 5.5.4 | 740 | 29 | 0 |
- Review round, 2026-10-30 (5648c93): the well-formedness test
checks each plan entry's style against the load's styles and that the
plan names a routed-experts container, for `prepared_ep_plan` too; a new
test checks every spec puts its plan on a config the load reads (a bad
style or an unread spelling fails them); `check_distributed_support`
drops a redundant local import. Pre-commit passes against main;
`tests/unit/torch/{models,utils}` pass on transformers 5.18 (777) and
5.5.4 (748, 29 skipped).
- Rebased on 2026-10-30 (fe16583) onto main 99c3834 through
#3131's rebase; no conflicts, this PR's own diff unchanged. Pre-commit
passes against main.
- 2026-10-30 (1bd24c0): rebased on #3134's review round; this PR's
own diff unchanged. Pre-commit passes against main;
`tests/unit/torch/models` and `test_moe_a2a.py` pass on transformers
5.18 (273) and 5.5.4 (255, 18 skipped).
- Review round, 2026-10-30 (3894689):
- Without transformers, the registry and plan checks run again; only the
tests that read the installed release skip. `test_ep_specs.py` runs 30
tests and skips 14 with transformers not importable, where it skipped
the whole module before.
- The FSDP2 way around expert parallelism also names an MTP head the
load cannot build.
- Pre-commit passes against main. `tests/unit/torch/models` and
`test_moe_a2a.py` pass on transformers 5.18 (273) and 5.5.4 (255, 18
skipped).
- Rebased on 2026-10-30 (da3aeca) onto #3134's rebase onto main
62a5c48, after #3131 merged: no conflicts, this PR's own diff
unchanged (range-diff). Pre-commit passes against main;
`tests/unit/torch/models` and `test_moe_a2a.py` pass on transformers
5.18 (312) and 5.5.4 (282, 19 skipped).
- Rebased on 2026-10-30 (9dd1d1e) onto main ff964eb after #3134
merged; this PR's own diff is unchanged (range-diff identical).
Pre-commit passes against main; `tests/unit/torch/models` and
`test_moe_a2a.py` pass on transformers 5.18 (312) and 5.5.4 (282, 19
skipped).
- Rebased on 2026-10-30 (85885c3) onto main 059588b, which moved
only by #3149, a CI change. This PR's own diff is unchanged (range-diff
identical). Pre-commit passes against main. `tests/unit/torch/models`
and `test_moe_a2a.py` pass on transformers 5.18 (312) and 5.5.4 (282, 19
skipped).
- Rebased on 2026-10-30 (8d2bd3c) onto main d464d30 (#3074 and
#3121 since 059588b). This PR's own diff is unchanged (range-diff
identical). Pre-commit passes against main. `tests/unit/torch/models`
and `test_moe_a2a.py` pass on transformers 5.18 (314) and 5.5.4 (283, 20
skipped).
- CodeRabbit nitpick, 2026-10-31 (41b8a6a): the test of the spec's
`prepare` hook now has it return a replacement config that carries the
plan over, and checks that the returned config is that one and holds the
plan. Pre-commit passes against main. `tests/unit/torch/models` and
`test_moe_a2a.py` pass on transformers 5.18 (314) and 5.5.4 (283, 20
skipped).
- CodeRabbit outside-diff item, 2026-10-31 (7c0be5b):
`apply_distributed_support` finds a plan the config already carries
before requiring the nested config the spec names. A config with a plan
of its own, at its root included, keeps it; with no plan anywhere, a
config without the nested config is still refused.
`test_a_plan_at_the_root_needs_no_nested_config` fails without the
change. Pre-commit passes against main. `tests/unit/torch/models` and
`test_moe_a2a.py` pass on transformers 5.18 (315) and 5.5.4 (283, 21
skipped).
- Rebased on 2026-10-31 (2b3dfc6) onto main e352cac, without
conflicts; its own diff is unchanged. Pre-commit passes against main.
`tests/unit/torch/models`, `test_moe_a2a.py` and
`test_model_load_utils.py` pass on transformers 5.18 (382) and 5.5.4
(339, 32 skipped).
- Rebased on 2026-11-01 (289b626) onto main cf38ec3, without
conflicts; its own diff is unchanged. Pre-commit passes against main.
`tests/unit/torch/models`, `test_moe_a2a.py` and
`test_model_load_utils.py` pass on transformers 5.18 (382) and 5.5.4
(339, 32 skipped).
- Review round, 2026-11-01 (ccab53a), the three follow-ups from the
code owner's review:
- `apply_distributed_support` refuses a `prepare` that returns `None`,
or a config with no plan where the load reads one, and names the model
type (b29d579). The new test's two cases fail without the check.
- The registered-plans check accepts `ep_router`, and its failure names
the styles the load keeps (202a805).
- `_spec_ep_plan`'s docstring covers a plan that ships with transformers
(ccab53a).
- Pre-commit passes against main. `tests/unit/torch/models`,
`test_moe_a2a.py` and `test_model_load_utils.py` pass on transformers
5.18 (384) and 5.5.4 (339, 34 skipped). The check changes nothing for
#2746: its tests with the check applied pass on 5.18 (588).
- Review round, 2026-11-01 (f9c898a): the registry test refuses a
spec declaring both a base and a prepared plan. Test-only. Pre-commit
passes against main. `test_ep_specs.py` passes on transformers 5.18 (47)
and 5.5.4 (32, 15 skipped).
- Rebased on 2026-11-01 (4cfbddf) onto main 51645db, without
conflicts; its own diff is unchanged. Pre-commit passes against main.
`tests/unit/torch/models`, `test_moe_a2a.py` and
`test_model_load_utils.py` pass on transformers 5.18 (384) and 5.5.4
(339, 34 skipped).

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ Additive: new functions, which
nothing calls until #2746 [3/3]. `_ep_config` and
`_num_experts_from_config` read the same two nested configs, in the same
order.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A. The feature's entry lands with #2754.
- Did you get Claude approval on this PR?: ❌ Not yet. Requested; #2746
as a whole was approved on 2026-10-28.

### Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models. Split out of
#2746 (OMNIML-6076), whose review history stays on #2746.







<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
<!-- coderabbit-ai-generated:v1 -->

## Summary by CodeRabbit

* **New Features**
* Added model-aware expert-parallel loading support for compatible
models, including models with nested configurations.
* Loading checks model support and Transformers version compatibility,
and explains when expert-parallel loading is unavailable, including
alternative loading options.
* Existing model-provided expert-parallel plans are preserved; supported
models can use their declared plans.
* Model-specific preparation is applied when available, helping ensure
the appropriate expert-parallel configuration is used.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

… load time (#3136)

> [!NOTE]
> **[2/3] of #2741, split on review:** #2741 grew past the 500-line
budget, so it is now three PRs, merged in order: #3131 [1/3] (the
helpers, merged), this PR (the expert-parallel loader), and #2741 [3/3]
(its post-load refusals).
>
> **Base:** `main`. Every PR this one builds on has merged: #3131 [1/3],
#2779 (on #3140), and #3067 and #2740 (on #3074). The helper branch
`shengliangx/hfep-base-6a-ep-loader`, which joined them until they
merged, is retired.

<details><summary>Series layout</summary>

- `main` (#3131 [1/3], #3140, #2779, #3074, #3067 and #2740 have merged)
  - **#3136 [2/3] expert-parallel loader (this PR)**
    - #2741 [3/3] the loader's post-load refusals
    - #2742 dense sharding → #2743 block-FP8 dequantization
- #2780 expert-parallel load routing (merges after #2741 [3/3]) → #2746
model specs
  - #3137 [2/3] on-the-fly expert decompression
- #2747 [3/3] decompressed checkpoints through the expert-parallel
loader (on #3136 and #3137)

</details>

### What does this PR do?

Type of change: New feature

`load_and_prepare_ep(ckpt_path, plan)` loads a model by an
expert-parallel `DistributionPlan`, with every rank both data-parallel
and expert-parallel.

- **Experts are sharded as they load.** transformers' own
expert-parallel load (`DistributedConfig(enable_expert_parallel=True)`)
builds only each rank's experts, over the `ep` submesh of an `(edp, ep)`
mesh.
- The model's `base_model_ep_plan` is swapped to the `torch_a2a` styles
for the load, so tokens travel to the rank that owns their expert.
- The config the model keeps gets the checkpoint's own plan back, so an
exported `config.json` names no ModelOpt-only style.
- The dense part is replicated, and every rank calibrates on its own
data (#2740).
- **The experts land on the EP group on transformers 5.18+ too.** Those
releases build a world-sized mesh of their own in `from_pretrained`; the
loader hands them the EP mesh for the load. Each `DistributedConfig`
argument goes where the release has a field for it.
- **Calibration's amax sync learns the layout from the model.** Each
module gets a `_parallel_state`: dense modules reduce over the whole
world, and routed-expert containers over their expert-data-parallel
replicas. The routed experts are found from the same
`base_model_ep_plan` that sharded them.
- **What is checked before anything loads:**
  - the plan and the process group;
- that every rank is bound to its GPU, and that no rank lacks one while
others have theirs (review fixes);
  - transformers 5.7 or later (`require_transformers`, #3131);
- the checkpoint's format: unquantized. MXFP4 is refused (review fix):
transformers dequantizes it into GPT-OSS's transposed expert layout
whatever the model's, so a standard-layout MoE would load its experts
transposed;
- that the plan styles routed experts, and that `ep_size` divides the
expert count, read from the config that holds the plan, as the
distribution plan's decision reads it (review fix).
- **The model keeps its EP mesh** for calibration
(`moe.expert_parallel_mesh`), and records the checkpoint tensors it has
no parameter for, for the export.
- **A rank that fails alone fails every rank** (review fix). Each rank
checks its own transformers, reads or is handed its own config, sees the
checkpoint where rank 0 made it local, and loads its own share, any of
which can fail on one rank alone. Each of those steps runs inside a
`_refused_together` block (#2779), whose failure-carrying `all_reduce`
raises the failure on every rank before the next collective:
  - the checks;
  - the checkpoint, made local through `ensure_local_checkpoint`;
- the load, through the routed experts' lookup. This rests on
`from_pretrained` running no collective, as transformers'
expert-parallel load does not on 5.7 through 5.18: before 5.16 it slices
each rank's experts into plain tensors, and from 5.16 it builds DTensors
locally (`run_check=False`, `src_data_rank=None`). The other ranks
finish the load and meet a rank that failed in it in the gather.

**Why a loader of its own.**
- **Not transformers' expert parallelism alone.** Its stock styles
(`ep_router`, `moe_tp_experts`) keep every rank on the same tokens. Each
rank runs all of them through its own experts, and an all-reduce over
the expert-parallel group sums the results. Calibration needs each rank
on its own data, so this load keeps transformers' expert sharding
(`grouped_gemm`) and swaps the routing for the `torch_a2a` styles. Those
send each token to the rank that owns its expert by all-to-all, around
the grouped GEMM that ModelOpt quantizes.
- **Not the meta/FSDP2 loader.** `parallel_load_and_prepare_fsdp2`
shards whole decoder layers, experts included, and gathers each layer's
weights on every rank in every forward. It cannot keep a rank's experts
resident and route tokens to them. Expert sharding also cannot start
from a meta-initialized model: transformers'
`GroupedGemmParallel.shard_tensor` slices materialized tensors. So this
load reads real weights through `from_pretrained` over the EP submesh.
- **The mesh-builder override.** From transformers 5.18,
`from_pretrained` drops its `device_mesh` argument once
`distributed_config` is given, and builds a world-sized mesh from the
config's sizes. Whenever the EP degree is below the world size, that
mesh shards the experts over every rank instead of the rank's EP group.
`_experts_sharded_over` hands transformers' builder the EP submesh this
load built, for the duration of the call only. Older releases honor
`device_mesh` and need no override.

### Usage

```python
from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan, size
from modelopt.torch.utils.moe import expert_parallel_mesh
from modelopt.torch.utils.plugins.model_load_utils import load_and_prepare_ep

# under torchrun, inside distributed_run()
plan = DistributionPlan(size(), ep_size=size())  # size() must divide the routed experts
model = load_and_prepare_ep(ckpt_path, plan, hf_config=config)
set_distribution_plan(model, plan)
ep_mesh = expert_parallel_mesh(model)  # for distributed_forward_loop(..., ep_mesh=ep_mesh)
```

### Design notes

Moved here from the code comments (review).

- **Transport.** The load swaps the model's `base_model_ep_plan` to the
`torch_a2a` styles. The router emits global expert ids unmasked. The
experts forward sends each token to the rank that owns its expert by an
NCCL all-to-all, runs the quantizable grouped GEMM, and sends the result
back by a second all-to-all, instead of every rank computing every token
and all-reducing. Every rank sees different tokens, hence per-rank
calibration data. Nemotron-H's logits-only gate has no `ep_router`
entry, so its `nemotron_experts_ep` also maps to `torch_a2a_experts`.
- **Meshes.** Two factorizations of the same world:
- `(edp, ep)` for the experts: sharded over `ep`, replicated over `edp`.
The `ep` submesh also drives `from_pretrained`'s expert sharding and the
transport.
  - `(dp,)` for the dense part, replicated over the whole world.

The all-to-all case is `edp == 1`, and the dense amax always reduces
over the whole world. `edp > 1` adds data parallelism across
expert-parallel groups, at the cost of per-rank memory, with no
concurrent `edp` and `ep` collectives.
- **Real weights.** Expert parallelism cannot be applied to a
meta-initialized model, because transformers'
`GroupedGemmParallel.shard_tensor` slices materialized tensors. So the
load reads real weights through `from_pretrained` over the `ep` submesh.
- **Release floor.** transformers 5.6 has the MoE expert-parallel
framework, but builds a model only with the experts implementations it
ships, never `torch_a2a`. A model whose expert-parallel support landed
later states its release in its spec.
- **No tensor parallelism.** transformers applies expert parallelism
through its tensor-parallel path: with `enable_expert_parallel=True`,
`model.tp_plan` returns the expert plan. `tp_plan="auto"` is only the
switch that makes `tp_size > 1` and so reaches
`apply_tensor_parallelism`. The swapped plan keeps expert entries only,
so attention and embeddings stay replicated.

Where `tp_plan` goes depends on the release: 5.17 and later refuse it
beside a `distributed_config`, and `DistributedConfig` has no `tp_plan`
field before 5.10, nor in 5.11 and 5.12. Passed alone on 5.17 and later,
it rebuilds the config without `enable_expert_parallel`. So the
dataclass fields decide: each argument goes on `DistributedConfig` where
it has a field for it, else to `from_pretrained`.
- **Collective-free load.** transformers' expert-parallel
`from_pretrained` runs no collective on 5.7 through 5.18. Before 5.16 it
slices each rank's experts into plain tensors, and from 5.16 it builds
DTensors locally (`run_check=False`, `src_data_rank=None`). So the other
ranks finish the load and meet a rank that failed in it in the guard's
gather.
- **Experts implementation.** `from_pretrained`'s
`experts_implementation` argument resolves non-FP8 implementations back
to `grouped_mm`, but the experts forward reads
`config._experts_implementation` when it runs, so the load sets it on
every (sub)config. Without that, the `torch_a2a_experts` style would
pass a fifth argument, the process group, to the grouped-mm forward.
transformers keeps it as `_experts_implementation_internal`, which a
serialized config leaves out.
- **Warmup.** The load creates each amax collective's NCCL communicator
in the same order on every rank, with a barrier around each, rather than
lazily mid-calibration. Lazy creation deadlocks when subgroups reach
their first collective in different phases at DP>1. There is no forward
warmup, since a variable per-group time would put the ranks out of step
first.
- **GPU binding.** Importing modelopt initializes CUDA wherever torchao
and triton are installed: torchao's mx_formats kernels call
`has_triton()`, which reads `torch.cuda.current_device()`. So in a
launch that binds no rank, every rank is already on cuda:0 when the load
runs, and `DeviceMesh`, which binds a rank only while CUDA is
uninitialized, leaves it there. On a gloo group, the load refuses such a
run on every rank: ranks on the same GPU, compared by UUID. An NCCL
launch in that state fails earlier, in `init_process_group`, where NCCL
refuses the duplicate GPU. Where CUDA is still uninitialized, the load
refuses ranks whose local rank names no GPU they see.

### Testing

- The expert-parallel load section of
`tests/unit/torch/utils/test_model_load_utils.py`, with the
`expert_parallel_mesh` cases in `test_moe_a2a.py` (CPU; in
`test_ep_loader.py` until 2026-10-31), covering:
  - the plan, group, floor and config refusals;
  - the Hub ID made local before the weights load;
  - the config's dtype and the class it builds through;
  - each `DistributedConfig` argument across the 5.6–5.19 field layouts;
  - the 5.18 mesh builder;
  - the restored plan and the saved config;
- the parallel-state seeding, `expert_parallel_mesh` and the
quantization-method read.
- **New:** `test_a_load_failing_on_one_rank_fails_every_rank` runs the
load on two real gloo ranks, with the meshes' collective stood in by a
bounded barrier. One rank fails alone at each guarded step:
  - an unreadable config;
  - transformers 5.6;
  - a GPTQ `hf_config`;
  - a checkpoint it cannot see;
  - out of memory in the load;
  - an unreadable weight file.

Every rank raises that failure. With the guards made rank-local, all six
rows fail, the other rank waiting at the barrier (about 8 s each,
bounded).
- `tests/gpu/torch/utils/test_expert_parallel_load.py`: loads a tiny
Qwen3-MoE expert-parallel with EP over every GPU and over half of them,
and checks each rank's share of the experts and the logits against the
whole model. Unchanged from #2741, where it passed on 2 and 4 GPUs. On
this split it ran at 96be72d (below).
- 2026-10-30: pre-commit, mypy included, passes against main.
`tests/unit/torch/utils` and the fused-experts tests pass on CPU at this
head, on transformers 5.18 and 5.5.4.
- Review fixes, 2026-10-30 (Claude's review): a plan in which some
routed-expert container has no `grouped_gemm` entry under it is refused
before anything loads, since `grouped_gemm` is the only expert-weight
style this load shards by (two new table rows fail without the check;
every plan transformers 5.18 ships and every ModelOpt spec's plan passes
it); the mxfp4 comments no longer name GPT-OSS; the docstring says to
pass `hf_config` on many ranks; `__all__` lists
`is_fsdp2_text_only_vlm`. Pre-commit, mypy included, passes against
main; the loader, model-loading and MoE unit tests pass on CPU on
transformers 5.18 and 5.5.4.
- Rebased 2026-10-30 onto main 99c3834, with #3131 rebased there
too; this PR's own diff is unchanged. Pre-commit, mypy included, passes
against main, and `tests/unit/torch/utils` passes on CPU at this head,
on transformers 5.18 and 5.5.4.
- Review fixes, 2026-10-30, round 2 (Claude's review): the load refuses
a plan whose routed-expert containers match no module of the loaded
model, inside its guard, so calibration cannot reduce the experts' amax
over ranks holding different experts; this refusal and its test move
here from #2741. It also refuses, before anything loads, a plan that
styles a weight under a routed-expert container by a style the load
drops; every `base_model_ep_plan` transformers 5.18 ships and every
ModelOpt spec plan passes the check. Both new tests fail without their
checks. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18 and 5.5.4.
- 2026-10-30: the no-module plan test's unplaced-keys stand-in returns a
list, as the function does (test only). The loader tests pass on CPU at
this head, on transformers 5.18 and 5.5.4.
- GPU run, 2026-10-30, at 96be72d (the guards around
`from_pretrained` on NCCL, until then checked on CPU gloo only):
`tests/gpu/torch/utils/test_expert_parallel_load.py` on NVIDIA GB200,
torch 2.11.0a0 (nv26.02), transformers 5.9.0, pytest with `-p
no:rerunfailures`.

  | GPUs | Result |
  |---|---|
  | 2 | 1 passed, 1 skipped (22 s) |
  | 4 | 2 passed (27 s) |
- Rebased 2026-10-30 onto the rebuilt
`shengliangx/hfep-base-6a-ep-loader`, which after #3131 merged holds
#2779 (on #3140, which moves `fsdp2_wrap` into
`modelopt/torch/utils/fsdp2_utils.py`), #3067 and #2740 on main
62a5c48. This PR's own diff is unchanged. Pre-commit, mypy included,
passes against main, and `tests/unit/torch/utils` passes on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU run above was taken
at 96be72d; the next entry is one on the rebased stack.
- GPU run, 2026-10-30, at #2743's head 9c917ae, which contains this
PR's head b726fb3 (#2742 adds the dense-sharded and
quantized-training cases to the file):
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

On 1 GPU every case skips, so it is no evidence. On 2 GPUs the four
cases with expert parallelism over half the GPUs skip; on 4 GPUs all 12
run: each rank's experts and the logits against the whole model, and FP8
and NVFP4 quantized training (scales and gradients), with the experts
over every GPU or half of them and the dense part replicated or
FSDP2-sharded.
- Review fixes, 2026-10-30 (meenchen's review bot):
- The expert count is read from the config that holds the plan
(4996a25). A new loader test on the decision test's nested config
fails without the fix.
- A load in which some ranks have a GPU of their own and others do not
is refused on every rank before either mesh (15e91d9). The bounded
two-rank test's new `no-gpu` row fails without the check.
- The load's helpers and the test file's transformers imports are
module-scope imports (3af30b9).
- The docstring says what the in-load failure guarantee rests on
(e251f6c).
- Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.
- Review fixes, 2026-10-30, round 2 (meenchen's and Claude's reviews):
- The load refuses an MXFP4 checkpoint (23feb60). On transformers
5.9.0 and 5.18.0, a tiny Qwen3-MoE loaded from a natural MXFP4
checkpoint got `gate_up_proj` `(4, 64, 192)` instead of `(4, 192, 64)`,
with no error. The config-refusal table's new `mxfp4` row fails without
the change.
- The GPU check is decided from one gather and names the GPU-less ranks
(98ee707). The docstring names the transformers releases the in-load
guarantee was checked on (1ba50c6).
- The quantization plugin's TP-linear converter no longer warns about a
mesh without TP-sharded Linears on a model recorded expert-parallel
(1491b66). Its new test fails without the change.
- Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4. So do the HF quantization plugin's tests.
- GPU run, 2026-10-30, at #2743's head 7192e05, which contains this
PR at e251f6c with the mixed-GPU gather:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

On 4 GPUs every case runs, the four with expert parallelism over half
the GPUs included, as at the run before. 1 GPU skips everything, so it
is no evidence.
- GPU run, 2026-10-30, at #2743's head 7c83230, which contains this
PR's head 1491b66:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

The same result as the run before, with the same test file, so it speaks
for the code underneath. The file loads an unquantized checkpoint on
transformers 5.18, so it cannot show the MXFP4 refusal or the block-FP8
floor of 5.8 firing: the config-refusal table's `mxfp4` row covers the
MXFP4 refusal, on CPU.
- Review fixes, 2026-10-30, round 3 (meenchen's review bot):
- The GPU check's predicate, which the load routing below defines
(#2779), now takes the device a rank is bound to once CUDA is
initialized, else the one its local rank numbers, as torch's
`DeviceMesh` does. The base is rebuilt on #2779 47ca6d1 for it. A
new bounded two-rank test (ad9559d) shows each process one GPU,
bound to `cuda:0` under local ranks 0 and 1: both ranks pass the check
and go on to the meshes. It fails with the old predicate. The `no-gpu`
row is unchanged.
- The tests build their model with the shared `model_with_decoder_at`
from `tests/_test_utils/torch/moe_models.py` (9a0f5e0).
- The style swap is defined once, before the checks, and the
dropped-style check and the swapped plan both read the styles it keeps
(3fc55b4).
- Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` and the HF plugin tests pass on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run
on this head.
- Rebased 2026-10-30 onto main 059588b, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
3b814a6, #3067 fbff3f3 and #2740 186efae; this PR's own
commits are unchanged.
- Review fix, 2026-10-30 (Claude's re-review): a run in which some rank
sees CUDA but is bound to no GPU yet is refused on every rank, telling
the caller to bind each rank to its GPU first, as `distributed_run()`
does (eb6c382). Before, such a rank was read as having no GPU and
the run was refused with "Run one process per GPU"; building the meshes
would have failed there anyway, as `DeviceMesh` binds an unbound rank by
its local rank. The one-GPU-per-process test runs that launch bound and
not yet bound; the not-yet-bound case fails without the check.
Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` and the HF plugin tests pass on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run
on this head.
- Review fixes, 2026-10-30 (Claude's review at eb6c382):
- The two-rank GPU-check tests pin, in every row, the CUDA state the
check reads, rather than read the host's (69179ec). The gloo workers
bind no GPU, so on a host with a visible one a row could be refused
before the step it exercises, depending on what earlier tests had
initialized. On a host with one visible GPU, a whole-file run at
eb6c382 failed the `no-gpu` row; on a host faked to show GPUs no
rank has bound, four rows failed. All pass now.
- The plan-check comment no longer says `packed_*` entries are dropped:
any other style under a routed-expert container is refused
(56353ef).
- No source behavior changed. Pre-commit, mypy included, passes against
main, and `tests/unit/torch/utils` and the HF plugin tests pass on CPU
at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- GPU run, 2026-10-30, at #2743's head 7540a2e, which contains this
PR at eb6c382: `tests/gpu/torch/utils/test_expert_parallel_load.py`
on transformers 5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

The same result as the runs before. On 4 GPUs every case runs, the four
with expert parallelism over half the GPUs included; 1 GPU skips
everything, so it is no evidence. Since that head, #3136 changed only
tests and one comment, so the result speaks for this head's code.
- Review fixes, 2026-10-30 (Claude's and meenchen's reviews at
56353ef):
- On a gloo group, the GPU check refuses a run in which two ranks are on
the same GPU (57548c6); an NCCL launch in that state fails in
`init_process_group`. Importing modelopt initializes CUDA wherever
torchao and triton are installed, so a launch that binds no rank has
every rank on cuda:0 at the load, and `DeviceMesh` leaves it there. The
not-yet-bound refusal is narrowed to a rank whose local rank names no
GPU it sees.
- The layout test covers one GPU each (bound, and not yet bound), every
GPU not yet bound, and every rank on cuda:0. The narrowing and the new
refusal each fail their cases without them.
- The single-process tests pin no GPU, through a `gloo_world_of_one`
override (b658f37).
- The docstring keeps the caller contract, and the comments a line or
two of why. The rationale is under Design notes (75bc2d6).
- At 56353ef on a host with one visible GPU, the unit file passed 48
of 48, and the two-rank test alone 7 of 7.
- Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` and the HF plugin tests pass on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run
on this head.
- Rebased 2026-10-30 onto main d464d30, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
8f9d877, #3067 d452768 and #2740 9a87d15; this PR's own
commits are unchanged. Pre-commit, mypy included, passes against main,
and `tests/unit/torch/utils` and the HF plugin tests pass on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4.
- GPU runs, 2026-10-30, at an earlier head with the same own diff
(#2743's 6f5391d, which contains this PR at 75bc2d6), on
transformers 5.18.0:
  - `tests/gpu/torch/utils/test_expert_parallel_load.py`:

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

On 4 GPUs every case runs, the four with expert parallelism over half
the GPUs included; 1 GPU skips everything, so it is no evidence.
- `tests/unit/torch/utils/test_ep_loader.py` on a host with one visible
GPU: 59 of 59 passed.
- The shared-GPU check, on one node with two processes seeing the same
two GPUs and neither bound, so that importing modelopt left both on
cuda:0:

    | Group | Result |
    |---|---|
| gloo | Both ranks refuse: "rank(s) [0, 1] are bound to the same one".
|
| nccl | `init_process_group` itself fails on both ranks ("Failed to
initialize NCCL communicator: invalid usage"), before the load. |
| nccl, one GPU per process (control) | No refusal; the load succeeds on
both ranks. |

So the check catches gloo groups; NCCL launches fail in
`init_process_group`.
- Tests moved 2026-10-31 (meenchen's review): the expert-parallel load's
cases from `test_ep_loader.py` to a section at the end of
`test_model_load_utils.py`, and the two `expert_parallel_mesh` cases to
`test_moe_a2a.py`, beside `moe.py`'s other tests (7747c5d). The
tests themselves are unchanged. Pre-commit, mypy included, passes
against main, and `tests/unit/torch/utils` collects and passes the same
cases as before on CPU, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-10-31 onto main e352cac, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
c78d2e5, #3067 adb9757 and #2740 dea9b79; this PR's own diff
is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` and the HF plugin tests pass on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4.
- GPU run, 2026-10-31, at #2743's head 939bf6e on main e352cac,
which contains this PR's head 930bb0f:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

The same 12 cases as before, now over the rebuilt base, which includes
main's fused-MoE GPTQ change. On 4 GPUs every case runs, the
dense-sharded ones included: each rank's experts and the logits against
the whole model, and both quantized formats' training. 1 GPU skips
everything, so it is no evidence.
- Rebased 2026-11-01 onto main cf38ec3, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
c42cb61, #3067 cb83366 and #2740 ff372fb; this PR's own diff
is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` and the HF plugin tests pass on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-01 onto main 51645db, where #3140 merged, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
1aee07f, #3067 495f0be and #2740 6a06a54; this PR's own diff
is unchanged. The tree is identical to the previous head, as #3140
merged as its own content. Pre-commit, mypy included, passes against
main.
- Rebased 2026-11-02 onto main 2d4e8d0, where #2779 and #3135
merged, on `shengliangx/hfep-base-6a-ep-loader` rebuilt there from #3067
e6e422d and #2740 5829a85 (#2779 is now part of main); this PR's
own diff is unchanged. Pre-commit, mypy included, passes against main,
and `tests/unit/torch/utils` and the HF plugin tests pass on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-02 onto main 4f08a12, where #3067 merged, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2740
5829a85 (#3067 is now part of main); this PR's own diff is
unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` and the HF plugin tests pass on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4.
- Retargeted 2026-11-02 to `main` and rebased onto main 653ef61,
where #2740 merged; the helper base `shengliangx/hfep-base-6a-ep-loader`
is retired, as every PR it joined has merged. This PR's own diff is
unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` and the HF plugin tests pass on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ Additive.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A. The feature's entry lands with #2754.
- Did you get Claude approval on this PR?: ❌ Requested. The loader was
approved as part of #2741 on 2026-10-28; the guards are new.

### Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models. The loader's
review history is on #2741.

























<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

<!-- coderabbit-ai-generated:v1 -->

## Summary by CodeRabbit

* **New Features**
* Added support for loading models for expert-parallel post-training
quantization, including validation of model configuration and
distributed setup.
* Added a way to identify whether a model uses expert-parallel
distribution and access its associated device mesh.
* **Bug Fixes**
* Avoided a misleading tensor-parallel warning when a model uses
expert-parallel distribution and has no tensor-parallel-sharded linear
layers.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…nd apply the model spec when deciding and loading it (#2780)

> [!NOTE]
> **Folded in #2746 [3/3] on 2026-11-03**, at the user's request: this
PR now also applies the model spec when deciding and loading
expert-parallel. #2746's commits are here unchanged (this PR's head is
#2746's last head, 22b6d23); #2746 keeps their review history, and
this PR covers its Jira ticket, OMNIML-6076, too.
>
> **Re-parented on 2026-10-30** onto #2741 [2/3] (#3136), the
expert-parallel loader's core, when #2741 was split. This PR also adopts
#2779's newer collective load checks; see the end of Testing.
>
> **Base:** `main`, since #2741 [2/3] (#3136) merged on 2026-11-03. The
diff below is this PR's alone. It merges after #2741 [3/3] (#2741),
whose refusals keep `load_pretrained`'s new expert-parallel route from
quantizing unsupported experts wrong.

<details><summary>Series layout</summary>

- `main`, with #3131 #2741 [1/3], #3134 #2746 [1/3], #3135 #2746 [2/3],
#2779, #3067, #2740 and #3136 #2741 [2/3] merged
  - #2741 [3/3] post-load refusals
  - #2742 dense sharding → #2743 block-FP8 dequantization
- **#2780 expert-parallel load routing, with the model spec applied
(this PR; #2746 [3/3] folded in)**
  - `hfep-base-11a-decompress` (#3137) → #2747 on-the-fly decompression
- `hfep-base-11b-kimi` (#2746's commits, now in this PR, + #2747) →
#3138 → #2748 → #2749
- `hfep-base-14-hf-ptq` (#2746's commits, now in this PR, + #2742 +
#2753 + #3133) → #3139 → #2754

</details>

### What does this PR do?

Type of change: New feature

#2779 routes `load_pretrained` by a `DistributionPlan`, but refuses a
plan that shards the routed experts: the expert-parallel loader (#2741)
comes later in the series. With it in, this PR adds that route, and
applies the model's spec (#3134), through the applier (#3135), when
deciding and loading it:

- **An expert-parallel plan loads through `load_and_prepare_ep`**,
handed the plan itself. Its arguments are checked with the FSDP2 load's,
in `_sharded_load_kwargs` inside #2779's collective "load checks" block;
`_load_sharded` then runs the load. That load replicates the dense part;
FSDP2-sharding it comes with #2742, a sibling of this PR. Its `freeze`
and `require_lm_head` options pass through, and so do the Hub options
(`token`, `cache_dir`, `local_files_only`), which its config read and
`from_pretrained` use. The sharded loads' refusals from #2779 cover it
too: a Nemotron-H checkpoint's MTP head, which `load_and_prepare_ep`
does not build, is refused rather than exported unquantized. Until #2742
is in, an expert-parallel plan that also shards the dense part is
refused before anything loads, naming the ways out: replicate it anyway,
or FSDP2-shard the whole model. A plan left to decide shards the dense
part when a whole copy of it and a rank's share of the experts would
exceed 60% of a GPU.
- **The decision refuses to leave an expert-parallel MoE whole on every
rank unasked**, naming the expert count the EP degree must divide and
every way out: shard its experts (an EP degree is offered only when one
divides both the experts and the world), shard the whole model with
FSDP2 (an unquantized checkpoint), load it in a single process, or ask
for the replication: `MODELOPT_EP_DEGREE=1` with `MODELOPT_FSDP_SIZE=1`,
or a replicated plan handed to `load_pretrained`, which skips the
decision. A replication both knobs ask for loads, with a note that every
rank holds the whole MoE. So #2737's refusals of routed experts the
decision can neither count nor shard by now advise
`MODELOPT_EP_DEGREE=1` with `MODELOPT_FSDP_SIZE=1` (`ep_size=1` with
`dense='replicate'` for a caller passing the arguments): an EP degree of
1 alone now meets the refusal above.
- **An MoE the expert-parallel load cannot take loads as a dense model
does**, whole on every rank or FSDP2-sharded, when the EP degree is left
to the decision or set to 1. The load may refuse it for the model's
spec, the transformers release it needs, its checkpoint format or its
plan. A rank-0 warning names the reason, and the fix where there is one
(an unquantized checkpoint, a later transformers); a replication both
knobs ask for gets only the "as asked" note. An EP degree above 1 asked
for such a model is still refused by the load.

**The model spec in the decision and the load** (folded in from #2746).
Without it, an MoE that only its spec makes expert-parallel, such as
DeepSeek-V3 before transformers 5.13, is decided as dense.

- **Deciding.** Unless the run rules expert parallelism out (a single
process, or an EP degree of 1), `resolve_distribution_plan` applies the
model's spec to a copy of the config before rank 0 builds the model on
meta.
- Every rank applies it inside a collective, so a rank that cannot fails
every rank.
- An expert-parallel plan returns the prepared copy. Any other plan
returns the config as read, which the export would otherwise write the
spec's plan into.
- Rank 0's knobs decide whether the spec applies, and rank 0 sends that
decision as it does the plan: a rank whose own environment skipped the
spec would load without the expert-parallel plan the others shard by,
and fail once the mesh is up.
- **Auto and the floors.** An EP degree left to auto resolves to 1 for a
model whose spec refuses expert parallelism, or that needs a later
transformers than the one installed: its spec's release or the
expert-parallel load's 5.7 floor, whichever is later.
- The load's floor holds for a model with no spec too
(`_spec_refuses_ep`).
- Only an explicit EP degree above 1, or a handed expert-parallel plan,
meets the refusal: `MODELOPT_FSDP_SIZE=world` alone plans FSDP2 for
GPT-OSS, and for Qwen3-MoE on transformers 5.5.
- **The left-whole refusal above sees spec plans.** At an EP degree of
1, which skips the spec, it still finds a spec-supplied plan's experts,
through the plan the spec would supply (or the one it states `prepare`
attaches), without applying anything.
- **An MoE whose spec rules expert parallelism out (GPT-OSS, Llama-4),
or whose transformers is below the release its expert parallelism needs,
loads as a dense model does** when the EP degree is left to auto or set
to 1. It loads whole on every rank, or FSDP2-sharded if it doesn't fit
one GPU, with the rank-0 warning above carrying the spec's reason or the
upgrade (`moe._spec_refusal`). An EP degree above 1 asked for such a
model is still refused, with the way around expert parallelism added
(`check_distributed_support`).
- **Loading.** The expert-parallel route applies the spec as it checks
the load's arguments, inside #2779's collective "load checks" block, so
a rank that cannot apply it fails every rank instead of leaving the
others in the load.
- A load handed the config that deciding prepared does not prepare it
again, since a model's preparation can patch its modeling code. It still
meets its spec's refusal and transformers floor
(`check_distributed_support`).
- A `config=` passed for the load with a plan handed in is the caller's
own, so an expert-parallel load prepares a copy of it, once.
- **Why the spec's `prepare` runs before deciding here, and where it
goes in #2748.** The decision builds the model on the meta device and
finds and counts its routed experts through the plan. For a model whose
plan only its spec supplies, the spec must therefore be applied first.
This PR applies the whole spec there, `prepare` included, and documents
the interim contract on `EPSpec.prepare`: it must leave a whole-model or
FSDP2 load as it would be without it, and may run more than once. No
spec in this PR sets `prepare`.
- #2748 adds the first specs that set it (Kimi-VL, Kimi-K2.5) and
replaces the contract. The decision runs only a pure-import step
(`EPSpec.make_importable`) and attaches the spec's stated plan
(`prepared_ep_plan`) with no `prepare` run. `prepare` then runs only
once the plan is decided expert-parallel: in a collective step after the
decision, or in the load after its own checks for a plan handed in.
- That puts `prepare` where `models/hf.py`'s `prepare_for_loading` hooks
already run: on the route that loads the model it prepares, never on one
that does not. The hooks are scoped to the load as a context; `prepare`
only runs on that route.

This closes the series' second phase: an MoE whose expert-parallel plan
ships with transformers, such as Qwen3-MoE, or that only its spec makes
expert-parallel, such as DeepSeek-V3, now loads expert-parallel through
`load_pretrained` and calibrates.

### Usage

```python
# torchrun --nproc_per_node 8 quantize.py
from modelopt.torch.utils.distributed import distributed_run
from modelopt.torch.utils.plugins.model_load_utils import load_pretrained

with distributed_run():
    model = load_pretrained("Qwen/Qwen3-30B-A3B")  # its plan ships with transformers
    # A model whose plan only its spec supplies loads the same way:
    # model = load_pretrained("/path/to/DeepSeek-V3-BF16")
    # ... mtq.quantize(model, ...), export_hf_checkpoint(model, ...)
```

Either model's routed experts are sharded over the 8 ranks. The
DeepSeek-V3 path is an unquantized (BF16) copy. The published
`deepseek-ai/DeepSeek-V3` checkpoint is FP8, which this PR refuses
before anything loads; the expert-parallel load dequantizes FP8
checkpoints from #2743 on.

### Testing

- `tests/unit/torch/utils/test_model_load_utils.py` (where #2737 moved
the decision's tests) gains the refusal to replicate an expert-parallel
MoE, when no degree divides its experts or an explicit EP degree of 1
leaves the dense part to auto, FSDP2 allowed instead, and the
replication both knobs ask for loading. Qwen3-MoE shows it, whose plan
ships with transformers 5.7.
`test_no_ep_degree_dividing_the_experts_leaves_them_whole` goes: the
refusal replaces it. #2737's two refusal tests now set
`MODELOPT_FSDP_SIZE=1` in their `set-to-1` rows and check the refusals
name it.
- `tests/unit/torch/utils/test_load_pretrained_distribution.py` (#2779)
gains the load's cases:
- an expert-parallel plan loads through `load_and_prepare_ep` with the
plan, records the plan, and goes through no other loader;
- the sharded loads' refusals (placement, model class, MTP head) apply
to it.
  - #2779's test that an expert-parallel plan is refused is gone.
- Mutation-checked: routing an expert-parallel plan elsewhere, dropping
the replicated-MoE refusal, and applying it to a sharded plan each fail
a test.
- `tests/gpu/torch/utils/test_load_pretrained_distribution.py` gains a
Qwen3-MoE case: left to decide, `load_pretrained` shards its experts,
each rank holding its contiguous share and the rest whole, and records
the expert-parallel plan. Not run locally; CI's GPU workflow runs it.
- Locally on this revision, the stack rebased on main c6fa613:
pre-commit over the whole PR against main, at every PR of the stack; and
at the stack's tip the full `tests/unit` on transformers 5.18 and 5.5
and on torch 2.8, and the hf_ptq example tests, each with the same
failures as main in that environment. CI runs the full unit suite.
- Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy
included, passes at this head over the whole change against main, and
the series' PRs merged together give a tree byte-identical to the former
stack's tip. This PR's own unit tests pass at this head on CPU (torch
2.14, transformers 5.18); CI runs the full unit suite.
- Review fixes, 2026-10-27: the refusal names no EP degree when the
config hides the expert count. Each behaviour fix has a test that fails
without it; pre-commit, mypy included, and this PR's own unit tests pass
at this head.
- Review fixes, 2026-10-27, round 2: a replication both distribution
knobs ask for loads, with a note; the refusal offers an EP degree only
when one divides both the experts and the world; the Hub options reach
the expert-parallel load. Each behaviour fix has a test that fails
without it; pre-commit, mypy included, and this PR's own unit tests pass
at this head.
- Rebased on 2026-10-27 onto main 7c3c868, after #2739 merged:
pre-commit, mypy included, and this PR's own unit tests pass at this
head, and the series' PRs merged together add no unit-test failure over
main (tests/unit/torch utils, export, quantization and models, on
transformers 5.18 and 5.5.4 and on torch 2.8).
- Review suggestions, 2026-10-27 evening: the expert-parallel load
passes the Hub options to the checkpoint download too.
- Re-cut on 2026-10-29 onto #2741 (was on #2743), with three follow-ups:
- the test hunks moved from `test_resolve_distribution_plan.py`, which
#2737 deleted, into `test_model_load_utils.py`;
- #2737's two "Set MODELOPT_EP_DEGREE=1" refusals name
`MODELOPT_FSDP_SIZE=1` too;
- the `load_pretrained` docstring says the expert-parallel route
replicates the dense part.

Pre-commit, mypy included, passes against main. On CPU (torch 2.14,
transformers 5.18 and 5.5.4), these pass at this head:
`test_model_load_utils.py`, `test_load_pretrained_distribution.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `test_distributed.py`.
- Review fix, 2026-10-29, from the re-review after the re-cut: on #2741
the expert-parallel load refuses a sharded dense part, so `_route` now
refuses an expert-parallel plan that also FSDP2-shards it, before
anything loads, naming the ways out. The routing test runs both dense
layouts; the refused one fails without the fix. Pre-commit passes
against main, and the unit-test files above pass on transformers 5.18
and 5.5.4.
- Review suggestion, 2026-10-29: the refusal says that sharding the
dense part under expert parallelism is not supported yet (message only).
Pre-commit passes against main.
- Re-parented on 2026-10-30 onto #2741 [2/3] (#3136); the former head is
pinned.
- Adapted to #2779's split of the sharded load: the expert-parallel
route's checks run in `_sharded_load_kwargs`, inside the collective
"load checks" block, and its load in `_load_sharded`.
- The Hub options reach the restructured loader's config read and
checkpoint step.
- Its tests in the deleted `test_load_pretrained_distribution.py` files
moved into `test_model_load_utils.py` (unit and GPU), where #2779 moved
the load tests.
- Review suggestion: the left-whole refusal's FSDP2 way out now says
"for an unquantized causal LM", as `_route`'s does (effe11d).

Pre-commit, mypy included, passes against main. CPU tests
(`test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py`;
torch 2.14):

| transformers | Passed | Skipped | Failed |
|---|---|---|---|
| 5.18 | 210 | 0 | 0 |
| 5.5.4 | 182 | 28 | 0 |
- Review suggestion, 2026-10-30 (e051283): when the expert-parallel
load refuses an option it does not take (such as `device_map`), it names
its way to a whole model, `MODELOPT_EP_DEGREE=1` with
`MODELOPT_FSDP_SIZE=1` or a replicated plan; the FSDP2 load's text is
unchanged. A test row checks it and fails without the change. Pre-commit
passes against main; `test_model_load_utils.py` and `test_ep_loader.py`
pass on transformers 5.18 (161) and 5.5.4 (139, 22 skipped).
- Rebased on 2026-10-30 (1bec6da) onto main 99c3834 through
#3136's rebase (4c59a6b); no conflicts, this PR's own diff
unchanged. Pre-commit passes against main; `test_model_load_utils.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models`
pass on transformers 5.18 (405) and 5.5.4 (376, 29 skipped).
- On GPUs, 2026-10-30: this PR's expert-parallel route, run as the
Qwen3-MoE row (`its-own-plan`) of
`test_an_expert_parallel_moe_loads_with_its_experts_sharded` at #2746
[3/3]'s ae518d6, passed. That version of the test also compares each
rank's logits with the whole model's. #2746's spec-plan row and the two
dense `load_pretrained` tests passed beside it, 4 tests in all. NVIDIA
GB200, torch 2.11.0a0, transformers 5.9.0, reruns disabled:

| GPUs | Passed | Failed |
|---|---|---|
| 2 | 4 | 0 |
| 4 | 4 | 0 |

- Rebased on 2026-10-30 (d78aaf8) onto #3136's restack (b726fb3)
after #3131 merged to main (62a5c48). The base under #3136 now
carries #2779's newer round, so 7 hunks in 3 files conflicted. All are
resolved by applying this PR's change to the new text:
  - `_load_sharded` takes the route beside #2779's new `_has_own_gpu`.
- `load_pretrained`'s docstring: the expert-parallel bullet goes before
#2779's reworded FSDP2 bullet, and the dropped "an `ep_size > 1` plan
raises" sentence leaves the rest of #2779's paragraph as it is.
- The tests keep #2779's new by-name, one-rank whole-load and
padded-auto cases beside this PR's. The `CHANGELOG.rst` line keeps this
PR's wording.
- Range-diff otherwise unchanged. Pre-commit passes against main.
`test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and
`tests/unit/torch/models` pass on transformers 5.18 (467) and 5.5.4
(425, 31 skipped).
- Review round, 2026-10-30 (50e64f0):
- The dtype refusal names FSDP2's `mp_policy` for the FSDP2 load only
(90c914a).
- A rank without a GPU of its own is refused before the expert-parallel
load too, as before the FSDP2 load. The remedy names the way to a whole
model for the route, from one helper shared with the placement refusal
(50e64f0).
- `test_a_sharded_load_refuses_what_it_cannot_honor` gains a `no_gpu`
row and an `mp_policy` assertion. Each fails without its change.
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models`
pass on transformers 5.18 (469) and 5.5.4 (427, 31 skipped).
- Review round, 2026-10-30 (3c353af):
- A whole-model load refuses the options only the expert-parallel load
takes (`freeze`, `require_lm_head`) before anything loads, naming that
load. Before, they reached `from_pretrained` and failed inside
transformers.
- `test_a_load_refuses_what_it_cannot_load` gains
`ep-option-on-a-whole-load`, which fails without the change.
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models`
pass on transformers 5.18 (470) and 5.5.4 (428, 31 skipped).
- Review round, 2026-10-30 (b1f5d8c):
- The left-whole refusal names the knobs as the caller set them:
`ep_size`/`dense` arguments, which override the environment knobs, or
the knobs (b1f5d8c). A new `explicit_degree_one_as_an_argument` row
fails without it.
- The docstring says the expert-parallel load freezes the model unless
`freeze=False`, unlike the other loads (1072e8f).
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models`
pass on transformers 5.18 (471) and 5.5.4 (428, 32 skipped).
- Review round, 2026-10-30 (4a1e020):
- The refusal of an MoE left whole on every rank is a `ValueError`, like
the decision's other refusals of what its inputs resolve to
(d0a63ff).
- A load handed an expert-parallel plan and no config reads the
checkpoint's config once, for its model type and for the load, as the
FSDP2 load already did (4a1e020). The EP-loader test asserts the
load gets that read.
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models`
pass on transformers 5.18 (471) and 5.5.4 (428, 32 skipped).
- Review round, 2026-10-30 (4040ea3):
- `freeze` or `require_lm_head` passed as `None` asks nothing, as other
options passed as `None` do, rather than reading as False and dropping
the LM head (83d9ce0).
- The expert-parallel route returns the model unfrozen unless
`freeze=True`, as the other routes do (69e56a9).
- A decided expert-parallel layout says it was decided when it refuses a
whole-model placement, as the FSDP2 route already did (4040ea3).
- Each change fails a test when reverted. Pre-commit passes against
main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py`
and `tests/unit/torch/models` pass on transformers 5.18 (472) and 5.5.4
(428, 33 skipped).
- Restack and review round, 2026-10-30 (8f9c8d9), onto #3136
1491b66. The rebase had no conflicts, and this PR's own commits are
unchanged; only the first commit's context moved.
- A decided expert-parallel layout now takes the load's own `freeze` and
`require_lm_head`, rather than refusing them as a whole-model placement
(1db4cc3). A new `its-own-options` row fails without the change.
- The docstring of `_refuse_a_rank_without_its_gpu` says what the check
adds beside #3136's GPU check in the expert-parallel load, which now
refuses a run where only some ranks have a GPU (8f9c8d9).
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models`
pass on transformers 5.18 (476) and 5.5.4 (430, 35 skipped).
- Review round, 2026-10-30 (d466147):
- A sharded load now refuses an option that only the other sharded load
takes by naming that load, rather than with the whole-model way out,
which would refuse the option again. `broadcast_chunk_size=None` asks
the FSDP2 load to broadcast every layer at once. The expert-parallel
load used to drop it silently and now refuses it as the FSDP2 load's
option. The two new rows of `test_a_load_refuses_what_it_cannot_load`
each fail without the change.
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models`
pass on transformers 5.18 (478) and 5.5.4 (432, 35 skipped).
- Restack and review round, 2026-10-30 (b692242), onto #3136
56353ef on main 059588b, after #2779 reworked the load checks:
- With an expert-parallel load to use it, the plan decision no longer
takes an EP degree left to auto as 1, so a decided MoE loads
expert-parallel. `test_a_load_it_decides_shards_the_experts` replaces
#2779's `test_a_load_it_decides_keeps_the_experts_whole`.
- This PR's refusals now use #2779's way out, `_whole_instead`, which
now takes the route. For a plan handed in, the refusal names a plan with
an EP degree of 1; for a decided plan, it also names the knobs.
- Placement options on a decided expert-parallel layout are now refused
by #2779's `measured` refusal, which this PR extends to that layout.
Only options that neither sharded load takes count as placement.
- Defaults now follow one rule, #2779's `_ASKS_NOTHING`. On the
expert-parallel load, `broadcast_chunk_size=None` is refused as an FSDP2
option, and `cpu_offload=False` asks for nothing.
- The refusal of the other sharded load's options now names a way
forward (e98619f).
- A launch with one GPU per process loads expert-parallel through
#2779's GPU predicate (5da80d4). That case fails under the
predicate's older reading.
- The changelog's backward-breaking entry now says an MoE loads
expert-parallel (b692242).
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models`
pass on transformers 5.18 (513) and 5.5.4 (463, 39 skipped).
- Review round, 2026-10-30 (1692141), onto #3136 75bc2d6:
- An EP degree left to the plan decision gives way to FSDP2 for the
whole model when the dense part, replicated beside a rank's share of the
experts, does not fit one GPU, instead of planning a layout the
expert-parallel load refuses (1692141). An EP degree the caller asks
for still stands and is refused by the load, naming the ways out.
`test_what_a_rank_holds_with_the_dense_part_replicated_must_fit` asserts
both outcomes, and its decided case fails without the change.
- The restack kept #3136's trimmed `load_and_prepare_ep` docstring and
comments, plus this PR's sentence on `hub_options`.
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models`
pass on transformers 5.18 (515) and 5.5.4 (463, 41 skipped).
- Restack and review round, 2026-10-30 (093c652), onto #3136
ab06860 on main d464d30:
- An EP degree the decision chose now also gives way to none when the
expert-parallel load would refuse the model, by its plan (the load's own
style checks, now one function, `_ep_plan_refusal`) or by its checkpoint
format. Such an MoE is not refused for being left whole either
(9cc4589).
`test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it`
fails without the change in both of its cases.
- The refusal of an expert-parallel plan that FSDP2-shards the dense
part says that an EP degree left to the decision gives way
(381219c).
- #2779's texts written before this load now name it: the auto_class
refusal, the no-group refusal and the docstring (093c652).
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models`
pass on transformers 5.18 (519) and 5.5.4 (464, 44 skipped).
- Review round, 2026-10-31 (1960174):
- An EP degree the decision chose now gives way for every reason the
config tells why the expert-parallel load would refuse the model: a
model spec ruling it out (Llama-4's, GPT-OSS's), a transformers below
the load's floor or the spec's, the checkpoint format, or the plan's
styles. The `its-spec` and `its-transformers` cases fail without the
change.
- Each give-way now says why on rank 0. The CHANGELOG names only the
MoEs the expert-parallel load takes as refused when left whole.
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models`
pass on transformers 5.18 (521) and 5.5.4 (464, 46 skipped).
- Restacked on 2026-10-31 (fac2c56) onto #3136 7747c5d, which
moved `test_ep_loader.py` into `test_model_load_utils.py`. This PR's one
change to that file, the Hub-options checks in
`test_the_weights_load_in_the_checkpoints_dtype_through_its_class`, now
sits in the moved section. Its own diff is otherwise unchanged.
Pre-commit passes against main. `test_model_load_utils.py`,
`test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers
5.18 (521) and 5.5.4 (464, 46 skipped).
- Review round, 2026-10-31 (3fc02b8):
- The expert-parallel load's transformers floor is one constant,
`_EP_MIN_TRANSFORMERS`, that the decision and the load both read
(6210606).
- The decision now resolves whether the expert-parallel load refuses an
MoE whatever its EP degree, so an EP degree of 1 the caller set no
longer meets the left-whole refusal for a model that load refuses
(3fc02b8). The give-way test runs each refusal with the degree on
auto and with `MODELOPT_EP_DEGREE=1`, and the `set-to-1` cases fail
without the change.
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers
5.18 (525) and 5.5.4 (464, 50 skipped).
- Review round, 2026-10-31 (3c40b57):
- `load_and_prepare_ep` refuses what the plan decision gives way for
(`_ep_load_refusal`), after its own floor and format checks. An EP
degree asked for, or a plan handed in, no longer loads a model whose
spec rules expert parallelism out. The load test's new GPT-OSS case
fails without the change.
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers
5.18 (526) and 5.5.4 (465, 50 skipped).
- Restacked on 2026-10-31 (c2d18e6) onto #3136 930bb0f, on the
new main e352cac and #2779 c78d2e5, with Claude's 04:08Z review
round folded in:
- #2779's texts meet the expert-parallel load. Its new
`init_quantized_weights` refusal names the expert-parallel route's own
way to a whole model, and a new test covers both a knob and a handed
plan; the plain FSDP2 wording fails it. The loading guide #2779 added
now describes the expert-parallel route, which this PR's docstring
paragraphs moved into (5455a0c). The decided EP degree is no longer
forced to 1 in `load_pretrained`.
- `_ep_load_refusal` returns the error itself, so a floor refusal is a
`RuntimeError` whichever floor it trips (eda4c88). A spec's floor
above the load's own is tested (608b0b8). Both cases fail without
their change.
- Pre-commit passes against main. `test_model_load_utils.py`,
`test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers
5.18 (536) and 5.5.4 (471, 54 skipped).
- Review round, 2026-10-31 (f7b11fe): the dense-sharding give-way
test also leaves the EP degree to the decision, where a caller asking
for FSDP2 makes the decided degree give way; it fails if only rank 0's
GPU memory may. Pre-commit passes against main. The same CPU tests pass
on transformers 5.18 (537) and 5.5.4 (471, 55 skipped).
- Restacked on 2026-11-01 (36863bc) onto #3136 4dd38a3, on main
cf38ec3 and #2779 c42cb61. #2779's new whole-model branch and
its two-rank test come in unchanged. The CHANGELOG entry takes #2779's
rewritten sentence, with this PR's MoE clause in place of the one saying
an MoE loads without expert parallelism. Its own diff is otherwise
unchanged. Pre-commit passes against main. `test_model_load_utils.py`,
`test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers
5.18 (538) and 5.5.4 (472, 55 skipped).
- Restacked on 2026-11-01 (8e19843) onto #3136 35f1ac0, on main
51645db, without conflicts; its own diff and tree are unchanged.
Pre-commit passes against main.
- Design change, 2026-11-01 (e20bd19): an MoE the expert-parallel
load cannot take now loads as a dense model does, with a rank-0 warning,
both when the EP degree is left to the decision and when it is set to 1
(before, only the first was noted, on rank 0, and the second loaded
silently). The give-way test asserts the warning for each reason the
load refuses, and fails without it. Pre-commit passes against main.
`test_model_load_utils.py`, `test_moe_a2a.py` and
`tests/unit/torch/models` pass on transformers 5.18 (538) and 5.5.4
(472, 55 skipped).
- Review round, 2026-11-01 (edf3b61): the dense-load warning skips a
replication both knobs ask for, which keeps its "as asked" note. The
give-way test asserts no warning there, and fails without the gate. The
guide promises a fix only where the load names one. Pre-commit passes
against main. The same CPU tests pass on transformers 5.18 (543) and
5.5.4 (472, 60 skipped).
- Restacked on 2026-11-02 (301f97d) on #3136 39f55f2, on main
2d4e8d0, where #2779 and #3135 merged. There were no conflicts, and
the PR's own diff is unchanged. Pre-commit passes against main. The same
CPU tests pass on transformers 5.18 (572) and 5.5.4 (486, 75 skipped).
- Restacked on 2026-11-02 (09322c1) on #3136 c0041d0, on main
4f08a12, where #3067 merged. There were no conflicts, and the PR's
own diff is unchanged. Pre-commit passes against main. The same CPU
tests pass on transformers 5.18 (572) and 5.5.4 (486, 75 skipped).
- Restacked on 2026-11-02 (7653ddf) on #3136 d87f5bd, which now
targets main 653ef61, where #2740 merged. There were no conflicts,
and the PR's own diff is unchanged. Pre-commit passes against main. The
same CPU tests pass on transformers 5.18 (572) and 5.5.4 (486, 75
skipped).
- Rebased on 2026-11-03 (e2d81c9) onto main a6cc167, where #3136
merged; this PR's base is now `main`. There were no conflicts, and the
PR's own diff is unchanged. Pre-commit passes against main. The same CPU
tests pass on transformers 5.18 (572) and 5.5.4 (486, 75 skipped).
- **From #2746, folded in on 2026-11-03.** #2746's Testing section keeps
its per-round record. Its tests and checks:
- `tests/unit/torch/utils/test_model_load_utils.py` (the decision's and
the loads' tests) gains the spec's cases:
- the spec applied before the decision, DeepSeek-V3's plan coming only
from its spec, and the returned config carrying it for an
expert-parallel plan only;
- no spec for a run that rules expert parallelism out, by argument or by
`MODELOPT_EP_DEGREE=1`;
- auto picks no EP degree for GPT-OSS, nor below a spec's floor or the
load's (DeepSeek-V3 on transformers 5.6.2); an explicit `ep_size=2`
meets the refusal;
- on two real gloo ranks, every rank applying the spec, or not, as rank
0 decided, whatever its own knob says;
- on two real gloo ranks, a rank that cannot apply the spec fails every
rank, when deciding and when loading a handed expert-parallel plan;
- the left-whole refusal for spec plans at an EP degree of 1; GPT-OSS,
and DeepSeek-V3 below the floor, loading as a dense model with the
warning, at an EP degree left to auto or set to 1;
- a decided expert-parallel load gets the spec-prepared config, prepared
once, not at all when handed, and on a copy when passed as `config=`
with a handed plan; GPT-OSS is refused with its config handed in,
without preparing it.
- Mutation-checked (earlier rounds): skipping the spec or applying it
without expert parallelism, ignoring the EP knob, letting each rank
decide the spec for itself, preparing a handed config twice or not at
all, ignoring the spec's plan at an EP degree of 1, letting auto pick an
EP degree for a model whose spec refuses expert parallelism or below a
floor, and offering an EP degree below the floor each fail a test.
- Checked locally, at the stack's tip: DeepSeek-V3, Nemotron-H and
Nemotron-H-Omni configs prepared by their specs and swapped to the
`torch_a2a` styles are written by the expert-parallel export with no
`base_model_ep_plan`, as their sources have none. There the export
writer itself puts back the source `config.json`'s values (#2753). The
loaded model's own config holds no spec plan only since b51123a (see
#2746's Testing).
- On GPUs, 2026-10-30, at ae518d6:
`test_an_expert_parallel_moe_loads_with_its_experts_sharded` in
`tests/gpu/torch/utils/test_model_load_utils.py` runs Qwen3-MoE
(`its-own-plan`) and DeepSeek-V3 saved with no plan of its own
(`its-spec-plan`) through the real `load_pretrained`, with the plan left
to decide. It checks each rank's expert slices and its own batch's
logits against the whole model. It and the two dense `load_pretrained`
tests passed on 2 NVIDIA GB200 GPUs (torch 2.11.0a0, transformers 5.9.0,
reruns disabled).
- Folded #2746 in on 2026-11-03 (22b6d23): this PR's branch
fast-forwarded to #2746's head, so no commit changed. Pre-commit passes
against main. The same CPU tests pass on transformers 5.18 (594) and
5.5.4 (495, 88 skipped).
- Review fix, 2026-11-03 (c9f3a2d): the refusal of a sharded load
inside `init_quantized_weights` gives each route its own reason. The
expert-parallel load calls `from_pretrained` with its own config and
device mesh, and the FSDP2 load's text is unchanged; the guide and a
test docstring follow. Pre-commit passes against main. The same CPU
tests pass on transformers 5.18 (594) and 5.5.4 (495, 88 skipped).
- Review fix, 2026-11-03 (275ce95): the Backward Breaking CHANGELOG
entry on how `load_pretrained` distributes a model is cut to two
sentences, what changed and how to keep loading it whole on every rank.
Pre-commit passes against main.
- Review fix, 2026-11-03 (0476abb): the guide's paragraph on calls
that place the model names the ways out the refusal gives, by route. For
an MoE whose experts the decision shards, that is `MODELOPT_EP_DEGREE=1`
with `MODELOPT_FSDP_SIZE=1`, or a plan with an EP degree of 1.
Pre-commit passes against main.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ Nothing released changes.
Against #2737 and #2779: an expert-parallel plan now loads, and a
decision to replicate an expert-parallel MoE on every rank now raises
unless both knobs ask for it.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ #2779's 0.49.0 Misc entry now names expert sharding.
- Did you get Claude approval on this PR?: This PR was approved on
2026-10-30 (effe11d), and #2746 on 2026-11-01 (04a197e), each
with its suggestions taken. Re-review requested for the folded PR.

### Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models; merges after
#2741 [3/3] (#2741). Split from #2774, which first routed every load in
one PR and whose non-expert-parallel half is now #2779: this half needs
the expert-parallel loader, so it lands after it. Applying the model
spec in the decision and the load, which this PR carried before the
series was reordered, moved to #2746 with the specs, and came back here
when #2746 was folded in on 2026-11-03.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Pretrained MoE models can load with routed experts sharded across
ranks under `torchrun`. Distribution plans also support FSDP2 sharding
or loading a full copy on each rank.
* **Bug Fixes**
* Loading now reports an error when a routed-expert model would
otherwise remain unsharded.
* **Documentation**
* Updated loading guidance and changelog details for distribution plans.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

… load, expert-parallel loads included (#3137)

> [!NOTE]
> **#2747 folded in (2026-11-03):** at the user's request to fold the
series backward, this PR now also carries #2747, the expert-parallel
loader's side of the decompression, with its commits as they were. #2747
is closed and keeps its review history; its ticket is this one,
OMNIML-6063.
>
> **Split out of #2747** at the reviewers' request (2026-10-30), and
joined with it again: the on-the-fly decompression,
`modelopt/torch/models/_onfly_decompress.py`, and the expert-parallel
load of a checkpoint its model spec decompresses.
>
> **Base:** `main`, which holds the expert-parallel loader (#3136) and
its routing with the model spec (#2780) since they merged.

<details><summary>Series layout</summary>

- `main` (#3130 [1/2], #3131, #3132, #3134, #3135, #3140, #2779, #3074,
#3067, #2740, #3136 and #2780 have merged)
  - #2741 [3/3] the loader's post-load refusals
  - #2742 dense sharding → #2743 block-FP8 dequantization
- **#3137 [2/2] on-the-fly expert decompression, expert-parallel loads
included (this PR; #2747 folded in)**
    - #3138 Kimi-VL and Kimi-K2.5 adapters (#2748 folded in)
      - #2749 [2/2] Kimi-K3 adapter

</details>

### What does this PR do?

Type of change: New feature

Expert-parallel PTQ loads a checkpoint with `from_pretrained` over the
expert-parallel mesh, so each rank builds only its own experts. A
checkpoint whose experts are compressed-tensors quantized (Kimi-K2.5's
INT4, Kimi-K3's MXFP4) cannot take that load as it is: transformers'
fused-experts conversion expects each expert's plain `.weight`.

- `install_onfly_decompress(hf_config)` rewrites the per-expert
converters of the model built from the config it marks.
- Each `.weight` source becomes the keys the checkpoint stores
(`weight_packed`, `weight_scale`, ...). A `DecompressExperts` op ahead
of the merge dequantizes each expert this rank owns to bf16, so PTQ runs
on the original compressed checkpoint with no decompressed copy on disk.
Experts other ranks own arrive as `None` and are skipped.
- It is format-agnostic: the `compressed_tensors` library does the math,
driven by the checkpoint's own `quantization_config`. The format is the
config group's own when it names one; the top level then says
"mixed-precision". A format compressed-tensors registers no compressor
for is refused.
- It handles a checkpoint whose experts are its only compressed weights,
in one config group that targets them. Several groups, or a group that
is not the experts, are refused; so is a `quantization_config` with no
groups, since there is nothing to decompress. Kimi-K2.5, K2.6 and
K2.7-Code (INT4) and Kimi-K3 (MXFP4) each carry one group, under
`text_config`.
- Stored params that disagree on how many experts they hold are refused,
as is a checkpoint key the decompression does not consume.
- The rewritten converter reverses as the one it rewrote: the
decompression is load-only, so a save (transformers' `save_pretrained`,
ModelOpt's export) writes the experts back under the checkpoint's own
`.weight` names (`w1`/`w3`/`w2` for Kimi-K3).
- The hooks are process-wide, but they rewrite only a marked config's
converters, and only a load of such a model keeps integer tensors
uncast: any other model loaded in the process keeps its stock converters
and casts.
- A second, differently compressed checkpoint in the same process is
refused; the same checkpoint again is a no-op.
- A config with nothing to decompress, or a raw dict, raises before
anything is patched, and without needing `compressed_tensors`. So does a
config that loads in a dtype other than bf16, the dtype the experts
decompress to, at any of its levels.
- It needs transformers 5.8: earlier releases' `Concatenate` looks up
the sibling keys the decompression has already consumed.
- It needs compressed-tensors 0.16, the first release whose compressors
decompress one weight by its scheme (`compression_param_names(scheme)`
and `decompress(state_dict, scheme)` as classmethods). An older one is
refused before anything is read, naming the release and the upgrade.
- `check_compressed_keys(hf_config, ckpt_dir)` reads the checkpoint's
stored tensor names and dtypes from its shard headers. Every compressed
weight must be an expert's, every expert's weight must be compressed,
every compressed weight must store the same names (the load pairs them
by position), all of them names the installed compressed-tensors gives
the scheme, and no compressed tensor may be a float bf16 cannot hold
exactly. The expert-parallel load runs it (below).
`stored_names(hf_config)` gives the names a marked config's checkpoint
stores, and refuses an unmarked config by name.
- CI's `unit` nox sessions install `compressed-tensors`, so the tests
against the real library run there. It goes in no extra: it pulls in
transformers, which CI's transformers-free lane must not get.

**Why not transformers' own compressed-tensors quantizer**
(`run_compressed=False`, which 5.18 deprecates for `dequantize=True`)?
It does not fit an expert-parallel PTQ load:
- Its fused-experts path, a `DecompressExperts` op its
`update_weight_conversions` prepends to the experts' converters, is
recent: transformers 5.5 has none. The Kimi-K2.5 adapter runs from 5.8.
- That op decompresses every expert it is handed into one buffer sized
for all of them, and has no case for the `None` the expert-parallel load
hands it for each expert another rank owns. This decompression shards at
load and decompresses only the rank's own experts.
- Before the load, the quantizer applies the checkpoint's
compressed-tensors scheme to every module its config groups target
(`apply_quantization_config`, `compress_model`). PTQ needs plain bf16
modules to place ModelOpt's quantizers on.

So the expert-parallel load deletes `quantization_config` on a
decompressed load (below), and the experts are decompressed here.

**The expert-parallel load of a decompressed checkpoint** (folded in
from #2747):

`load_and_prepare_ep` takes a checkpoint whose model spec decompresses
its experts as they load (above), which its format check refused before.

- **Format check.** A config the decompressor marked passes the check;
an unmarked compressed-tensors checkpoint is still refused.
- **No competing quantizer.** On the marked path the load deletes
`quantization_config` from its copy of the config, at the top level and
on `text_config` and `llm_config`: every level `_ckpt_quant_method`
reads. So transformers builds no compressed-tensors quantizer to compete
with the decompressor.
- Kimi-K2.5's config class copies its text config's
`quantization_config` to the top level, and from transformers 5.13 that
quantizer claims the experts' converters itself.
- It is deleted rather than set to `None`: before 5.11, transformers
takes any `quantization_config` attribute, `None` too, for a quantized
checkpoint.
  - The caller's config keeps its own.
- **Stored tensors checked before anything loads
(`check_compressed_keys`).** A class target such as `Linear` does not
say what was compressed, and the load drops the config that would load
any other compressed weight. So every weight stored compressed must be
an expert's, and every expert's weight stored compressed. No compressed
tensor may be a float that bf16 cannot hold exactly.
- **No compressed tensor carried unread.** The export copies the
checkpoint tensors the model has no parameter for as stored, and this
load dropped the config that would read a compressed one. So once the
load knows those tensors, it refuses any that is compressed: an MTP
layer's compressed experts, say, which pass the stored-tensor check but
sit in no layer transformers builds. Uncompressed ones are carried as
before. So are compressed tensors in a safetensors file the load never
opens, such as a standalone `mtp.safetensors`, which the export copies
as stored; a plain such file is carried.
- **Collective checks.** The format check and the deletion run inside
the loader's collective checks (#3136), and so does the stored-tensor
check, right after the checkpoint is made local. The unplaced-tensor
refusal runs inside the load's own collective block. So a refusal on one
rank fails every rank.

**Non-expert-parallel loads.** A spec's `prepare` installs the
decompressor and marks the config, and it runs only on expert-parallel
routes:
- only an expert-parallel plan runs it, on every rank; the plan decision
reads the plan the spec states instead (since #2748, now in #3138);
- an EP degree of 1 and a single process skip it;
- the FSDP2 load refuses a quantized checkpoint and a VLM before it
loads.

So a marked config reaches `from_pretrained` only through this load.
What outlives a `prepare` in the process is inert for other loads: the
conversion-mapping hook rewrites only a marked config's converters.

### Usage

Internal to the Kimi adapters (#3138, #2749), whose `prepare` installs
the decompressor; the expert-parallel load does the rest.

### Testing

- `tests/unit/torch/models/test_onfly_decompress.py` (49 cases) covers:
- the rewritten converters, a converter source that is no weight refused
at the rewrite, and a stored param without its weight refused at the
load;
- packed experts decompressing and fusing like an uncompressed
checkpoint, through the loader's own steps;
- only owned experts reaching the compressor, and stored params
disagreeing on the expert count refused;
- integer tensors keeping their dtype through transformers' own copy in
a decompressing load, and a later load of an unmarked model casting as
before (a real `from_pretrained`);
- installing once, a second checkpoint compressed differently refused,
and a config with nothing to decompress, a raw dict, or a load in a
dtype other than bf16 at any level, refused;
- the format taken from the one group, and a format with no compressor
refused;
  - an unmarked config refused by name when its stored names are read;
- the stored-tensor check over ten layouts, among them scales stored for
other experts in equal number and a name the installed
compressed-tensors does not give the scheme;
  - a VLM marked at either level;
- real `from_pretrained` loads of a tiny compressed Mixtral, INT4 and
MXFP4, with the real library: the experts load exactly, save back as the
checkpoint's own `w1`/`w3`/`w2` weights, and export under those names.
These and the other cases with the real library run in CI's unit
sessions; elsewhere they skip, saying what they need.
- Skips below transformers 5.8, the release the decompression requires.
- Locally at the round-1 head:
  - pre-commit against main passes, mypy included;
- on CPU, torch 2.14: 87 passed on transformers 5.18, and 45 passed, 3
skipped on 5.5.4. That is `test_onfly_decompress.py` plus the
expert-parallel loader's `test_ep_loader.py`.
- 2026-10-30: round 2 (an unmarked config refused by name, the group's
own format and a format with no compressor refused, a non-bf16 load
refused, docstrings compressed), rebased onto #3136 4c59a6b (main
99c3834); `tests/unit/torch/models/test_onfly_decompress.py` +
`tests/unit/torch/utils/test_ep_loader.py` +
`test_hf_checkpoint_utils.py` on CPU, torch 2.14: 146 passed on
transformers 5.18, and 94 passed, 9 skipped on 5.5.4; pre-commit against
main passes, mypy included. Each fix's test fails without it.
- 2026-10-30: rebased onto #3136 b726fb3, restacked after #3131
merged, with no conflicts and the own diff unchanged. Pre-commit against
main 62a5c48 passes, mypy included. On CPU, torch 2.14,
`test_onfly_decompress.py`, `test_ep_loader.py`,
`test_unplaced_source_keys.py`, `test_hf_checkpoint_utils.py` and
`test_model_load_utils.py` pass: 301 on transformers 5.18, and 237
passed, 21 skipped on 5.5.4.
- 2026-10-30: round 3 (cjluo-nv's review): the converter reverses as the
one it rewrote, compressed weights must store the same names, integers
stay uncast only in a decompressing load, bf16 at every config level,
module-level imports. Each fix's test fails without it. Pre-commit
against main passes, mypy included. On CPU, torch 2.14,
`test_onfly_decompress.py`, `test_ep_loader.py`,
`test_unplaced_source_keys.py`, `test_hf_checkpoint_utils.py`,
`test_model_load_utils.py` and `tests/unit/torch/export` pass on
transformers 5.18 and 5.5.4, apart from seven `test_export_diffusers.py`
cases that fail the same way at the base in this environment.
- 2026-10-30: rebased onto #3136 e251f6c (its module-scope imports
and mixed-GPU refusal) with no conflicts and the own diff unchanged.
Pre-commit against main passes; the same five test files pass on
transformers 5.18 (308) and 5.5.4 (238 passed, 22 skipped).
- 2026-10-30: Claude's two suggestions on round 3 taken: a stored name
the installed compressed-tensors does not give the scheme is refused,
and so is a converter source that is no weight, where it is rewritten.
Each test fails without its change. Pre-commit against main passes; the
five test files pass on transformers 5.18 (309) and 5.5.4 (238 passed,
22 skipped).
- 2026-10-30: juhi's blocker fixed. From transformers 5.14 an
expert-parallel load shards each expert's tensors through
`DtensorShardOperation.shard_tensor`, which cast the packed integers to
bf16 (a 2-GPU Kimi-K3 load on 5.18 failed in `DecompressExperts`). The
keep-integers patch now sits on `spawn_materialize`, where every read
starts; a new test shards an int32 tensor through transformers' own
`DtensorShardOperation` on a one-rank CPU mesh and fails without the
fix. Claude's suggestions taken: experts' weights only rewritten, an
expert's bias passes the stored-name check, the transformers floor
checked first, wording. Rebased onto #3136 1491b66. Pre-commit
against main passes; `test_onfly_decompress.py`, `test_ep_loader.py`,
`test_unplaced_source_keys.py`, `test_hf_checkpoint_utils.py` and
`test_model_load_utils.py` pass on transformers 5.18 (313) and 5.5.4
(239 passed, 22 skipped), with compressed-tensors 0.19.0. The 2-GPU run
of the compressed expert-parallel load on 5.18 (#2747's GPU case) is
pending on our cluster.
- 2026-10-30: On two GPUs, with transformers 5.18.0 and
compressed-tensors 0.19.0,
`tests/gpu/torch/utils/test_expert_parallel_load.py::test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own[int4]`
and `[mxfp4]` passed at #2747 e2e7486 (on this PR's 0d7866e),
loading each case at EP2 through the `spawn_materialize` fix; on one GPU
both skip, as they need two. Since then this PR gained only the
module-scope test import (b56c3e7) and shorter docstrings
(48b922c), with no behaviour change; source is now 495 added lines,
within the ~500 budget.
- Review fixes carried from #2747:
  - `_SUB_CONFIG_ATTRS` is imported at module level;
- the unit sessions install `compressed-tensors` on every torch, since
the torch 2.8 and 2.9 legs it was kept out of are gone.
- 2026-10-30: rebased onto main 059588b through #3136 584c8b4,
whose base now sits on main with #3134 merged. No conflicts, and the own
diff is unchanged. Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils` pass on CPU,
torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-10-30: restacked onto #3136 56353ef (main 059588b). From
Claude's review of #2747: `check_compressed_keys` returns the `tensor
name -> file` map it read, for the loader to reuse (7a56a68);
nothing else changed. Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils` pass on CPU,
torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-10-31: restacked onto main d464d30 through #3136 ab06860.
No conflicts; the own diff is unchanged. Pre-commit against main passes,
mypy included; `tests/unit/torch/models` and `tests/unit/torch/utils`
pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-10-31: at 43f08b5, the decompressor ran in #2747's
compressed-expert GPU cases, `[int4]` and `[mxfp4]`, which passed on 2
GPUs and on 4 GPUs with transformers 5.18.0 and compressed-tensors
0.19.0, in a run of the whole GPU file at #2749 63c8267, which
contains it.
- 2026-10-31: restacked onto main e352cac through #3136 930bb0f,
which moved the expert-parallel loader tests into
`test_model_load_utils.py`. No conflicts; the own diff is unchanged.
Pre-commit against main passes, mypy included; `tests/unit/torch/models`
and `tests/unit/torch/utils` pass on CPU, torch 2.14, on transformers
5.18 and 5.5.4.
- 2026-10-31: at 9deea3b, over the rebuilt base, the decompressor
ran in #2747's compressed-expert GPU cases, `[int4]` and `[mxfp4]`,
which passed on 2 GPUs and on 4 GPUs with transformers 5.18.0 and
compressed-tensors 0.19.0, in a run of the whole GPU file at #2749
a60c01c, which contains it.
- 2026-10-31, at c6b427a: on compressed-tensors 0.13.0 the
decompressor's setup failed with `TypeError: 'property' object is not
callable`. Before 0.16 a compressor states its parameter names as a
property (to 0.14) or not at all (0.15), and decompresses through an
instance built from the whole config. An older compressed-tensors is now
refused up front, naming 0.16 and the upgrade (321ba9a,
c6b427a), and the tests that need the real library skip below 0.16.
`test_a_compressed_tensors_before_per_scheme_compressors_is_refused_by_its_release`
fails without the check, on a stand-in compressor of 0.13's form.
Pre-commit against main passes, mypy included; `tests/unit/torch/models`
and `tests/unit/torch/utils` pass on CPU, torch 2.14, on transformers
5.18 and 5.5.4. `test_onfly_decompress.py` against each
compressed-tensors release, on CPU, torch 2.14, transformers 5.18:

  | compressed-tensors | Passed | Skipped |
  |---|---|---|
  | 0.19.0, 0.16.0 | 54 | 0 |
  | 0.15.0.1, 0.13.0 | 45 | 9 |
- 2026-10-31, at 81bcef7, from Claude's review: the loader hook
reads the decompressing flag first and looks at a source's dtype only in
a decompressing load, so any other load passes through it untouched;
`test_integer_checkpoint_tensors_keep_their_dtype_on_load` hands it a
source with no dtype outside such a load, which raised before. The
test's reverse name mapper is imported at module scope. Pre-commit
against main passes, mypy included; `tests/unit/torch/models` and
`tests/unit/torch/utils` pass on CPU, torch 2.14, on transformers 5.18
and 5.5.4.
- 2026-11-01, at e72a810: restacked onto main cf38ec3 through
#3136 4dd38a3. No conflicts; the own diff is unchanged. Pre-commit
against main passes, mypy included; `test_model_load_utils.py` and
`test_onfly_decompress.py`, run in one process, pass on CPU, torch 2.14,
on transformers 5.18 and 5.5.4.
- 2026-11-01, at df2de9d: restacked onto main 51645db through
#3136 35f1ac0; the tree is unchanged. Pre-commit against main
passes, mypy included; `test_model_load_utils.py` and
`test_onfly_decompress.py`, run in one process, pass on CPU, torch 2.14,
on transformers 5.18 and 5.5.4.
- 2026-11-02, at fcd6c08: restacked onto main 2d4e8d0 through
#3136 39f55f2. Main now holds #3135, whose `_SUB_CONFIG_ATTRS` in
`utils/moe.py` is the one this PR added, word for word, so the PR no
longer changes `moe.py`; the rest of the own diff is unchanged.
Pre-commit against main passes, mypy included; `tests/unit/torch/models`
and `tests/unit/torch/utils`, run in one process, pass on CPU, torch
2.14, on transformers 5.18 and 5.5.4.
- 2026-11-02, at f410618: restacked onto main 4f08a12, which now
holds #3067, through #3136 c0041d0. No conflicts; the own diff is
unchanged. Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils`, run in one
process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-11-02, at 59edbb4: restacked onto main 653ef61, which now
holds #2740, through #3136 d87f5bd, which now targets main. No
conflicts; the own diff is unchanged. Pre-commit against main passes,
mypy included; `tests/unit/torch/models` and `tests/unit/torch/utils`,
run in one process, pass on CPU, torch 2.14, on transformers 5.18 and
5.5.4.
- 2026-11-02, at c9dcaa9: rebased onto main bc5adc6, without
#3136, and retargeted to main: its three files import nothing #3136
adds. No conflicts; the own diff is unchanged. Pre-commit against main
passes, mypy included; `tests/unit/torch/models` and
`tests/unit/torch/utils`, run in one process, pass on CPU, torch 2.14,
on transformers 5.18 and 5.5.4, and without transformers.
- 2026-11-02, at ac48ad2, from CodeRabbit's review: every
in-function import in `_onfly_decompress.py` and its test now names its
reason in a comment: optional dependencies, and `core_model_loading`
being new in transformers 5. Comments only; pre-commit against main
passes, and `test_onfly_decompress.py` passes on CPU, torch 2.14, on
transformers 5.18, and skips on 5.5.4 and without transformers, as
before.
- 2026-11-03, at a39cdb5: #2747 folded in. This PR's branch was
fast-forwarded to #2747's head, so its commits are #2747's as they were,
on main a6cc167. Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils`, run in one
process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4. On
GPU, the whole GPU file passed at #2749 7cfcd54, which contains this
code, with transformers 5.18.0 and compressed-tensors 0.20.1: on 2 GPUs,
5 passed and `[ep-half]` skipped, as it needs 4; on 4 GPUs, all 6
passed, in the file's order and again in reverse.
- 2026-11-03, at 24bc5b6: the history was linearized onto main
a6cc167, replaying this PR's and #2747's commits in order without
the merge from main, so DCO no longer counts main's squash. The tree is
identical to a39cdb5's, every commit is signed off, and pre-commit
against main passes.
- 2026-11-03, at c3cfcc6, from CodeRabbit's review: the loader tests
import `_onfly_decompress`, packaging's `Version` and transformers'
version at module scope, since nothing makes them optional there;
`core_model_loading` stays imported in its test, after the version
check, with its reason beside it. Tests only; pre-commit against main
passes, and `test_model_load_utils.py` and
`test_unplaced_source_keys.py` pass on CPU, torch 2.14, on transformers
5.18 and 5.5.4.
- 2026-11-03, at 25a4cd8, from Claude's review:
`to_native_backbone_config` moved to #3138, where the Kimi adapters call
it, with its tests; this PR keeps the VLM case of the either-level load
test. Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils`, run in one
process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-11-03, at eff418c, from cjluo-nv's review: a decompressing
load refuses compressed tensors in the safetensors files it never opens,
which the export copies as stored (ea3f31e; the test fails without
it), and the real round trips give each expert scales of its own and
store packed weights and scales across two shards in opposite expert
order (eff418c), passing on transformers 5.8.0 and 5.18.0.
Pre-commit against main passes, mypy included; `tests/unit/torch/models`
and `tests/unit/torch/utils`, run in one process, pass on CPU, torch
2.14, on transformers 5.18 and 5.5.4, and without transformers.
- 2026-11-03, at 2d8208d, from Claude's review: a decompressing load
reads the checkpoint's shard headers once, as the unplaced-key check
takes the map `check_compressed_keys` read. Pre-commit against main
passes, mypy included; `tests/unit/torch/models` and
`tests/unit/torch/utils`, run in one process, pass on CPU, torch 2.14,
on transformers 5.18 and 5.5.4, and without transformers.
- 2026-11-03, at 86452a8, from kevalmorabia97's review: the tensor
map reaches `unplaced_source_keys` by keyword and only when the load
read one, so a two-argument replacement records its keys again (a new
test pins both forms; it fails with the positional call), and the tests
that write over a safetensors file they loaded clone its tensors first.
Pre-commit against main passes, mypy included;
`tests/unit/torch/export`, `tests/unit/torch/models` and
`tests/unit/torch/utils`, run in one process, pass on CPU, torch 2.14,
on transformers 5.18 and 5.5.4, apart from seven
`test_export_diffusers.py` cases that fail the same way on main here,
and without transformers.
- 2026-11-03, at 8feea6b: rebased onto main 7d44246, which now
holds #2780. Three conflicts, in imports and context only: the loader's
import list keeps the names #2780 uses beside this PR's `config_levels`,
the local-checkpoint step keeps #2780's hub options ahead of this PR's
stored-tensor check, and the loader tests import `_onfly_decompress`
beside `get_spec`. The tests' `_NO_SPEC` constant is no longer added
here, as #2780 added it word for word. The rest of the own diff is
unchanged. Pre-commit against main passes, mypy included;
`tests/unit/torch/export`, `tests/unit/torch/models` and
`tests/unit/torch/utils`, run in one process, pass on CPU, torch 2.14,
on transformers 5.18 and 5.5.4, apart from seven
`test_export_diffusers.py` cases that fail the same way on main here,
and `tests/unit/torch` passes without transformers.

- Earlier history and GPU results: see #2747.

**Tests folded in from #2747:**

- `tests/unit/torch/utils/test_model_load_utils.py`: the
compressed-checkpoint case loads from a small safetensors checkpoint,
which stores int32 packed weights and bf16 scales as the real ones do.
It runs the load to `from_pretrained`, whose config has no
`quantization_config` at any level, `llm_config`'s included, while the
caller's config keeps its own.
  - The same checkpoint unmarked is refused.
- One that also compresses an attention projection is refused before
`from_pretrained`.
- A compressed tensor the built model has no parameter for (an MTP
layer's expert) is refused once the weights load; an uncompressed one
gets past the check.
- The stored-name check, run read-only on the tensor names of the real
Kimi-K2.5, K2.6, K2.7-Code and K3 checkpoints (their index and shard
headers): all four pass. Reading the stored names, and with them the
compressed experts' dtypes, takes, per rank:

  | Checkpoints | Names | Names and dtypes |
  |---|---|---|
| Kimi-K2.5, K2.6, K2.7-Code, K3 | 0.5 s to 2.2 s (the most: K3's
497,220 names, read cold) | not timed |
| Synthetic, K3's 247,500 compressed tensors in 10 shards (CPU) | 0.5 s
| 1.0 s |

- On GPUs, the Kimi-K3 run (#2749) loads its MXFP4 experts through this
path.
- `tests/unit/torch/utils/test_unplaced_source_keys.py`: a real
decompressed `from_pretrained` on CPU (a tiny int4 pack-quantized
Qwen3-MoE through the real decompressor): every built expert's stored
tensors are placed through the recorded converters, and only a layer
past `num_hidden_layers` stays unplaced.
-
`tests/gpu/torch/utils/test_expert_parallel_load.py::test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own[int4|mxfp4]`
(2 GPUs): a tiny Qwen3-MoE whose experts are compressed (INT4, MXFP4,
through compressed-tensors) loads at EP2 through the decompressor; each
rank's experts must equal the dense model's exactly, every stored tensor
must be placed, and the logits must match the dense model loaded whole.
- #2747's dated testing history, its 2- and 4-GPU runs included, stays
on #2747.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ Additive; the decompressor is
installed only by a model's `prepare` hook, and only a config a model
spec marked takes the new load path.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A.
`compressed_tensors` is imported only for compressed checkpoints, which
already need it; only CI's unit sessions install it, for the tests.
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A. The feature's entry lands with #2754.
- Did you get Claude approval on this PR?: ✅ Approved on 2026-11-03 with
#2747 folded in (at eff418c). Before the fold, this PR's round 1 was
approved on 2026-10-30.

### Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models. #2747's split
is now two PRs:
- [1/2] is #3130, merged;
- [2/2] is this PR, with #2747 folded back in.

Jira: OMNIML-6063, which #2747 shared.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
<!-- coderabbit-ai-generated:v1 -->

## Summary by CodeRabbit

* **New Features**
* Load supported compressed Mixture-of-Experts checkpoints through
Transformers, with expert weights decompressed during loading and
converted to bfloat16.
* Expert-parallel loading distributes the relevant expert weights across
workers while preserving integer tensor data types.
* Checkpoint validation catches unsupported compression settings,
inconsistent stored weights, and compressed tensors outside supported
model layers.
* Requires Transformers 5.8 or newer and compressed-tensors 0.16 or
newer.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…o each rank (#2742)

> [!NOTE]
> **Re-parented on 2026-10-30 onto #3136**, #2741 [2/3] (the
expert-parallel loader). #2741 was split into three PRs on review: #3131
[1/3], #3136 [2/3] and #2741 [3/3], its post-load refusals. This PR
needs only the loader, so it sits on [2/3] and no longer waits for
[3/3]. The series first became a tree on 2026-10-27 (this PR was
[10/21]).
>
> **Base:** `main`, where #3136 and #2780 have merged (2026-11-03). The
diff below is this PR's alone. #2780 merged first, so this PR carries
what the two need together: it lifts #2780's refusal of an
expert-parallel plan with a sharded dense part, and the plan decision's
give-way to whole-model FSDP2.
>
> **Re-parent:** this PR's diff is unchanged, except that its dense-part
load re-anchors around #3136's load guard. `from_pretrained`, with this
PR's warmup and dense-meta contexts, runs inside the guard. The dense
part's FSDP2 + DCP load, which runs collectives, comes after it. The
PR's original 37 commits are condensed into one; the review-round
commits are kept.

<details><summary>Series layout</summary>

- `main` (#3131 [1/3], #3136 [2/3] and #2780 have merged)
  - #2741 [3/3] the loader's post-load refusals
  - **#2742 dense sharding (this PR)**
    - #2743 block-FP8 dequantization

</details>

### What does this PR do?

Type of change: New feature

#2741 replicates the dense (non-expert) part of an expert-parallel model
on every rank. That does not fit when the dense part alone is bigger
than a GPU, which the plan from #2737 already recognizes
(`dense="fsdp"`).

- **The dense part is FSDP2-sharded over every rank when the plan says
so** (`dense="fsdp"`, with the experts sharded over every rank too), and
nothing beside the plan decides it:
  - it is left on meta during `from_pretrained`;
- each decoder layer is then `fully_shard`-ed around its experts, which
stay expert-parallel. The blocks come from Hugging Face's own
`_no_split_modules` (`_fsdp2_wrap_modules`), looked up on the language
backbone rather than on a VLM wrapper that would also list its vision
blocks. A VLM whose language model shows no blocks (no
`_no_split_modules` match, no detected decoder layers) is refused,
naming `MODELOPT_FSDP_SIZE=1`, instead of wrapped by the VLM's own
`_no_split_modules`, which can name vision blocks;
- each rank loads only its own shards of the dense weights from the
checkpoint through DCP (#2739).

No rank ever holds the whole dense part. A dense param the loader
already placed as a replicated DTensor is re-registered as a plain local
parameter first, so FSDP2 shards it like any other.
- **A block-FP8 checkpoint meets the loader's format refusal here,
whatever its dense layout.** This load does not dequantize block FP8;
#2743 adds that, with the rule that keeps a block-FP8 dense part
replicated.
- **A dense part FSDP2-sharded within several expert-data-parallel
replicas is refused before loading.** The no-gather export cannot write
that layout, so `load_and_prepare_ep` refuses a plan with `dense="fsdp"`
and `edp > 1` before it looks at anything, rather than after the load
and calibration it would otherwise run. The error says to set
`MODELOPT_EP_DEGREE` to the world size, or `MODELOPT_FSDP_SIZE=1`. A
replicated dense part at `edp > 1` stays supported. The decision knows
the layout too: left to decide, it replicates such a dense part and
warns that it may run out of memory, naming the same way out; asked to
shard it, it refuses with the same advice.
- **The load-time allocator warmup is sized to the real per-rank
footprint.** transformers' `caching_allocator_warmup` otherwise reserves
about the whole device on every rank, and torch keeps that reservation,
starving the all-to-all's NCCL buffers. Now:
  - expert params count their 1/ep share;
  - dense params left on meta are not counted at all.

The replacement forwards whatever else transformers passes the hook, and
a hook whose leading parameters are not `(model, accelerator_device_map,
hf_quantizer)` keeps the stock sizing. Warmup sizes measured while
building this (moved here from the source):

  | Model, layout | Sized before | Sized with this PR |
  |---|---|---|
| Nemotron-3-Ultra-550B, EP32 | ~275 GiB per rank (experts at full size)
| ~91 GiB per rank |
| Kimi-K3, EP64, dense part sharded | 185.8 GiB on a 184.31 GiB card
(dense part on meta counted at full size) | dense part on meta not
counted |

A dense parameter off the model's dtype (an fp32 router) stays
replicated, as in `fsdp2_wrap`, and the dense meta load refuses a
transformers without the `get_device` hook it patches.

- **A rank that fails before the DCP load fails every rank.**
`_meta_dcp_load` (main's, from #2739; this PR's dense load is its only
caller) runs its per-rank preparation (the checkpoint's metadata read,
the mapping checks, `to_empty`) through the failure-carrying
`all_reduce` before `dcp.load`, so one rank failing there no longer
leaves the others waiting in the load. Main's FSDP2 load
(`parallel_load_and_prepare_fsdp2`) fills its shards through
`set_model_state_dict`, not this function, so the fix reaches only this
PR's dense load.
- **Every alias of a tied parameter is loaded.** The dense load names
both aliases of a tied pair, so a tied checkpoint that stores only
`lm_head.weight` fills both, as `from_pretrained` does.

### Why shard the dense part, and why separate helpers

**A model that needs it.** Kimi-K3 at EP64 on GB200s (184.31 GiB each).
With the dense part replicated, a rank holds the dense part whole and
its 1/64 share of the routed experts. transformers' load warmup sized
that at 185.8 GiB, more than the card before any activation. With the
dense part sharded, 79.2 GiB of weights stayed resident per rank. The
NVFP4 run over 16 GB200 nodes finished in about 10 minutes at a peak of
155 GiB per GPU, on an earlier build of this series.

**What it costs where it is not needed.** Qwen3-30B-A3B at EP16 on 4
nodes × 4 GB200s (transformers 5.9.0, torch 2.11). Both arms ran
`hf_ptq.py --recipe general/ptq/nvfp4_default-kv_fp8 --batch_size 2
--calib_size 512 --calib_seq 2048 --skip_generate` with
`MODELOPT_EP_DEGREE=16`, back to back in one allocation:

| Per rank (max over the 16 ranks) | Replicated (`MODELOPT_FSDP_SIZE=1`)
| FSDP2-sharded (`MODELOPT_FSDP_SIZE=16`) |
|---|---|---|
| Peak allocated after load | 6.35 GiB | 4.64 GiB |
| Peak allocated after calibration | 8.41 GiB | 6.83 GiB |
| Peak reserved after calibration | 17.36 GiB | 16.90 GiB |
| Load time | 6.2 s | 5.8 s |
| Calibration time | 20.7 s | 21.2 s |
| Export time | 4.2 s | 4.5 s |

Sharding saves 1.7 GiB of resident dense weights per rank and costs
about 0.5 s (2%) of calibration in the dense all-gathers. This model
does not need it: its dense part is about 1.8 GiB. With
`MODELOPT_FSDP_SIZE` unset, `auto` chose `replicate` and reproduced the
replicated arm to within 0.01 GiB and 0.1 s. The two layouts export the
same checkpoint (see Testing).

**The 60% boundary is a heuristic.** `auto` compares what a rank would
hold with the dense part replicated (the dense part whole, plus its 1/ep
share of the routed experts) with 60% of its GPU's memory. The other 40%
is left for calibration's activations and the communication buffers, the
all-to-all's among them. The two measured models sit far on either side
of the line: Kimi-K3 just over 100% of the card before any activation,
and Qwen3-30B-A3B at a few percent. No sweep across models and GPU sizes
has placed the line itself. Where it picks wrong, `MODELOPT_FSDP_SIZE=1`
replicates the dense part and `MODELOPT_FSDP_SIZE=<world size>` shards
it.

**Why separate helpers, rather than extending the existing paths:**
- **`parallel_load_and_prepare_fsdp2`** builds the model from its
config, FSDP2-shards every decoder layer with its experts, and has each
layer's owner read the layer and broadcast it to every rank. Under
expert parallelism the routed experts are not FSDP2-sharded.
`from_pretrained` places each rank's share of them on the
expert-parallel mesh through transformers' tensor-parallel plan, and the
all-to-all forward routes tokens to them. The hybrid keeps
`from_pretrained` for the experts and shards only the rest; reusing the
FSDP2 loader would mean re-implementing that placement in it.
- **`fsdp2_wrap`** wraps each detected decoder layer whole, and the root
by default, and refuses a model its detection misses. The hybrid has
different wrap units:
- each block is wrapped around its experts, which stay out of the FSDP2
groups;
  - the root stays replicated;
- blocks come from the language backbone's `_no_split_modules`, so a
VLM's vision blocks stay replicated.

Folding that into `fsdp2_wrap` would change the wrap units and the
accepted models of the existing FSDP2 load, so its callers keep it as it
is.
- **Accelerate offload** cannot combine with this load. transformers
refuses a `device_map` beside the tensor-parallel placement the
expert-parallel load requests: "`tp_plan` and `device_map` are mutually
exclusive" on 5.5.4, and the same refusal on 5.18. Offload would also
keep a whole copy of the dense part per rank in host memory and move it
to the GPU for every calibration forward. Sharding holds 1/world of it
and all-gathers one block at a time.

### Usage

```python
plan = DistributionPlan(size(), ep_size=size(), dense="fsdp")
model = load_and_prepare_ep(ckpt_path, plan, hf_config=config)
```

### Testing

- In `tests/unit/torch/utils/test_model_load_utils.py` (first written as
`test_ep_dense_sharding.py`), on CPU under gloo:
- telling expert params from dense ones through the EP plan, under any
prefix, with routers kept replicated;
- a block-FP8 checkpoint with a sharded dense part refused for its
format;
- a dense part sharded within several expert-data-parallel replicas
refused before anything loads (`test_ep_loader.py`);
- dense weights resolving to meta during the load, with the device hook
and the warmup's sizing restored even after an error;
- each decoder layer sharded around its experts, and the blocks to wrap:
from `_no_split_modules` on the language backbone, falling back to
decoder-layer detection, and refused for a model without blocks and a
VLM whose language model shows none (`test_meta_dcp_load.py`);
  - a replicated DTensor param sharded like a plain one;
  - the shard-then-load order.
- In `tests/unit/torch/utils/test_model_load_utils.py` (first written as
`test_ep_load_warmup.py`):
- the per-rank sizing of experts, meta dense weights, buffers and
quantized elements;
  - transformers reserving exactly that footprint;
  - falling back to the stock sizing.
- Mutation-checked: each test fails on at least one deliberate break of
the code.
- The expert-matcher cases are one table of 7 rows, and other duplicates
become cases of the tests that cover them. Each drop is checked by
mutation.
- The decision replicating, or refusing, a dense part within several
expert-parallel groups; an fp32 router left replicated, with every layer
still unsharding; the missing-`get_device` refusal.
- Locally on this revision, the stack rebased on main c6fa613:
pre-commit over the whole PR against main, at every PR of the stack; and
at the stack's tip the full `tests/unit` on transformers 5.18 and 5.5
and on torch 2.8, and the hf_ptq example tests, each with the same
failures as main in that environment. CI runs the full unit suite.
- GPU results at a named head are under "GPU runs, 2026-10-30" at the
end of this section; the other GPU results here were measured on earlier
revisions of this branch.
- `tests/gpu/torch/utils/test_expert_parallel_load.py` now also runs
with the dense part sharded, at EP over every GPU (sharded within
several EP groups, the loader refuses it): each rank's experts and the
logits match the whole model. At 8192a75 the whole file passed: 12
on 4 GPUs; 8 on 2 GPUs, where the 4 cases with EP over half the GPUs
skip.
- Also in `tests/gpu/torch/utils/test_expert_parallel_load.py`, QAT
under expert parallelism: an FP8 or NVFP4 model calibrated
expert-parallel must reach the whole model's scales, with its experts
over every rank or half of them and the dense part replicated or
FSDP2-sharded, and its gradients through the fake quantizers must match.
These tests came with the fused-experts quantization (#2745) before the
series was reordered; they land here, the first PR with an FSDP2-sharded
dense part. At 8192a75, where the whole file passed:

  | Dense layers, EP span | 2 GPUs: FP8, NVFP4 | 4 GPUs: FP8, NVFP4 |
  |---|---|---|
  | replicated, all ranks | pass, pass | pass, pass |
  | replicated, half the ranks | skipped | pass, pass |
  | FSDP2-sharded, all ranks | pass, pass | pass, pass |

- End to end with the export fixes in #2750, the sharded and replicated
dense layouts export the same checkpoint:
  - tiny Qwen3-MoE, EP=2 on 2 GPUs: byte-identical for NVFP4 and FP8;
- Qwen3-30B-A3B, EP16 over 4 nodes: byte-identical weights and weight
scales for NVFP4 and FP8, with only calibration-derived input scales
differing.

Both were measured before #3074's expert weight-scale fix, when both
layouts inflated expert scales the same way: the agreement shows the two
layouts match, not that those scales were right.
- Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy
included, passes at this head over the whole change against main, and
the series' PRs merged together give a tree byte-identical to the former
stack's tip. This PR's own unit tests pass at this head on CPU (torch
2.14, transformers 5.18); CI runs the full unit suite.
- Review fixes, 2026-10-27: a dense sharding that finds blocks but
shards none is refused, pointing at `MODELOPT_FSDP_SIZE=1`; a TP-planned
param keeps its share before the warmup's meta skip, and the TP-plan
lookup tries `transformers.distributed` first. Each behaviour fix has a
test that fails without it; pre-commit, mypy included, and this PR's own
unit tests pass at this head.
- Review fixes, 2026-10-27, round 2: the warmup replacement forwards
whatever transformers passes the hook and leaves a hook of another
signature alone; replicating a dense part within several expert-parallel
groups warns that it may run out of memory, with the way out; the
dense-sharding comments and progress line say where it wraps and that it
shards over every rank. Each behaviour fix has a test that fails without
it; pre-commit, mypy included, and this PR's own unit tests pass at this
head.
- Rebased on 2026-10-27 onto main 7c3c868, after #2739 merged:
pre-commit, mypy included, and this PR's own unit tests pass at this
head, and the series' PRs merged together add no unit-test failure over
main (tests/unit/torch utils, export, quantization and models, on
transformers 5.18 and 5.5.4 and on torch 2.8).
- Rebased over #2741's last round, which prints the loader's notes
through `print_rank_0` and no longer imports `is_master`: this PR's
dense-sharding notes import it themselves.
- Review suggestion, 2026-10-28: two of the dense-sharding notes print
through `print_rank_0`; `is_master` stays where it keeps the
meta-parameter listing off the other ranks.
- Review suggestion, 2026-10-28, second round: the EP-aware
warmup-sizing note prints through `print_rank_0` too; this PR's
MoE-helper import joins #2737's module-level one.
- Review suggestion, 2026-10-28, third round: the warmup sizing credits
a TP-planned param `1/ep_degree`, its share of the ep submesh
`from_pretrained` places it over, rather than `1/world_size`. The test's
world now holds two ep submeshes and fails with the old divisor.
- Declined review suggestion, 2026-10-28, third round: alias-safe
re-registration of replicated parameters tied inside one decoder block,
which no model this load supports has; the embedding-to-`lm_head` tie
sits at the root, which is not wrapped.
- Review suggestion, 2026-10-28, fourth round: the warmup's docstring
says only that no expert param is in `_tp_plan` under the EP load, which
is what the sizing relies on.
- Review suggestion, 2026-10-28, fifth round: the expert-param matcher
is the dot-padded containment test alone, which the exact and suffix
matches were special cases of.
- Review suggestions, 2026-10-28, sixth round: the expert matcher's
docstring states its one containment test, and the meta-placement
context and the sharded dense load each have their own summary.
- Re-cut on 2026-10-29 (see the note at the top): rebased onto the
re-cut #2741 with no source conflict. The edits of the deleted
`test_resolve_distribution_plan.py` (the env-knob table, and the dense
part the expert-parallel load cannot shard) moved to
`test_model_load_utils.py`. Pre-commit, mypy included, passes at this
head over the whole change against main. On CPU at this head (torch
2.14), this PR's tests with `test_ep_loader.py`, `test_meta_dcp_load.py`
and `test_model_load_utils.py` pass on transformers 5.18 and 5.5.4.
- Review fixes, 2026-10-29 (cjluo-nv's review bot):
- A rank failing in `_meta_dcp_load`'s preparation now fails every rank
before `dcp.load`. A bounded two-rank test fails rank 1 either in its
metadata read or with `to_empty` out of memory; with the fix reverted,
rank 0 enters the load. (Corrected 2026-10-30: main's FSDP2 load does
not call `_meta_dcp_load`, so the fix reaches only this PR's dense
load.)
- The dense load names both aliases of a tied parameter. A tied tiny
Qwen3-MoE storing only `lm_head.weight` failed through this caller on
the CPU before the fix, and now loads both aliases. The load also takes
its device from the mesh.
- Imports moved to module level, and the helper docs compressed; the
benchmark numbers moved above.
- Pre-commit, mypy included, passes against main.
`tests/unit/torch/utils` passes on the CPU at this head on transformers
5.18 and 5.5.4. GPU runs of that head did not happen; see 8192a75
below.
- Re-parented 2026-10-30 onto #3136 (#2741 [2/3]):
- `from_pretrained` re-anchored inside the load guard; the dense part's
load after it.
- Claude's three suggestions taken: the GPU test checks every rank for a
missing gradient before its collective comparisons; `_meta_dcp_load`
hands its prepared state over as a named tuple;
`replace_layer_number_by_wildcard` is imported from either transformers
location.
  - Pre-commit, mypy included, passes against main.
- `tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18 and 5.5.4.
  - The GPU tests did not run at that head; see 8192a75 below.
- Review fixes, 2026-10-30 (Claude's re-review): with the dense part
sharded, a load in which no dense param came back on meta is refused,
since every rank would then have loaded the whole dense part
(`test_a_dense_part_from_pretrained_did_not_leave_on_meta_is_refused`
fails without the check); the EP-aware warmup sizing warns when it falls
back to transformers' own sizing (the fallback test asserts the warning
and fails without it). Rebased with #3136 onto main 99c3834; the
rest of this PR's diff is unchanged. Pre-commit, mypy included, passes
against main, and `tests/unit/torch/utils` passes on CPU at this head,
on transformers 5.18 and 5.5.4.
- Review fixes, 2026-10-30 (cjluo-nv's review bot):
- `_fsdp2_wrap_modules`' two imports stay in the function, as
`fsdp2_wrap`'s does, with a comment on why (6388dc4; superseded
below, they are now at module scope): the export and quantization
packages sit above this module, and the export package imports it.
Moving them would not cycle today: with both at module scope, every
entry order imports on transformers 5.18 and 5.5.4 and on torch 2.8, and
`import modelopt.torch` already loads both modules. So the reason is
layering, not cost.
- The description now carries the Qwen3-30B-A3B costs, the Kimi-K3 case,
the 60% rule and the reasons for separate helpers. It also corrects two
claims: main's FSDP2 load does not call `_meta_dcp_load`, and block-FP8
handling lands in #2743.
- Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18 and 5.5.4.
- Rebased onto #3136 96be72d, which now refuses an expert-parallel
plan matching no module, or styling an expert weight other than
`grouped_gemm`. The dense-meta refusal test's unplaced-keys stand-in
returns a list, as the function does.
- The expert-parallel GPU tests ran at 8192a75 on 2 and 4 GPUs
(below).
- Review fixes, 2026-10-30, round 2 (cjluo-nv's review bot):
- `_fsdp2_wrap_modules`' two imports, and `fsdp2_wrap`'s, are at module
scope (c2f1bf9). Each of `model_load_utils`, `modelopt.torch.export`
and `modelopt.torch.quantization`, imported first in a fresh
interpreter, loads cleanly on transformers 5.18, 5.9.0 and 5.5.4.
- The EP-aware warmup sizing also warns when transformers' hook takes
other leading parameters, the way a transformers upgrade would reach its
fallback; the hook-signature test asserts the warning and fails without
it. The refusal of a dense part sharded within several
expert-data-parallel replicas also names `MODELOPT_EP_DEGREE=1`, which
keeps the dense part sharded by dropping expert parallelism
(201e2a5).
- Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- GPU runs, 2026-10-30, at 8192a75: NVIDIA GB200, torch 2.11.0a0
(nv26.02), transformers 5.9.0, pytest with `-p no:rerunfailures`.

  | Test file | 2 GPUs | 4 GPUs |
  |---|---|---|
| `tests/gpu/torch/utils/test_expert_parallel_load.py` (EP logits,
backward, FP8/NVFP4) | 8 passed, 4 skipped (25 s) | 12 passed (94 s) |
| `tests/gpu/torch/utils/test_model_load_utils.py` (main's FSDP2 load) |
6 passed (26 s) | not run |
| `tests/gpu/torch/export/test_fsdp2_export.py` (main's FSDP2 load) | 39
passed (147 s) | not run |

These results were taken at 8192a75. Up to aa85e91 this PR's
diff since then was the module-scope imports, the hook-signature
warning, the `MODELOPT_EP_DEGREE=1` remedy text and two docstrings; the
2026-10-30 rebase below then changed the base under it, and the file was
re-run on the rebased stack (the last entry here).
- Review suggestions, 2026-10-30 (Claude's review at 201e2a5,
approved): two docstrings no longer over-claim. `_fsdp2_wrap_modules`
leaves vision blocks unwrapped only for a VLM
`get_language_model_from_vl` recognizes, and the dense meta load leaves
buffers alone unless they are named `.weight`/`.bias` (aa85e91).
Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Review fix, 2026-10-30 (cjluo-nv's review bot):
`test_ep_dense_sharding.py` and `test_ep_load_warmup.py`, which covered
`model_load_utils` alone, moved into `test_model_load_utils.py` with
their fixtures and parametrization, as CONTRIBUTING.md's naming rule
asks (004de2a). No test was added or dropped:
`tests/unit/torch/utils` gives the same counts on CPU on transformers
5.18, 5.9.0 and 5.5.4, and pre-commit, mypy included, passes against
main.
- Rebased 2026-10-30 with #3136 onto the rebuilt
`shengliangx/hfep-base-6a-ep-loader`, which after #3131 merged holds
#2779 (on #3140, which moves `fsdp2_wrap` into
`modelopt/torch/utils/fsdp2_utils.py`), #3067 and #2740 on main
62a5c48. With `fsdp2_wrap` in its own module, this PR's paragraph in
its docstring moved to `fsdp2_utils.py`, `fsdp2_wrap` keeps its local
quantization import there (`fsdp2_utils` imports no Hugging Face or
quantization package), and `model_load_utils` imports
`_off_dtype_params` and `fully_shard` for the dense sharding itself. The
rest of this PR's diff is unchanged. Pre-commit, mypy included, passes
against main, and `tests/unit/torch/utils` passes on CPU at this head,
on transformers 5.18, 5.9.0 and 5.5.4.
- GPU run, 2026-10-30, at #2743's head 9c917ae, which contains this
PR's head f991fa5:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0, the test file unchanged since the run at 8192a75.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

On 1 GPU every case skips, so it is no evidence. On 2 GPUs the four
cases with expert parallelism over half the GPUs skip; on 4 GPUs all 12
run: each rank's experts and the logits against the whole model, and FP8
and NVFP4 quantized training (scales and gradients), with the experts
over every GPU or half of them and the dense part replicated or
FSDP2-sharded.
- Rebased 2026-10-30 onto #3136 e251f6c (meenchen's review round),
whose module-scope imports replace the load's local import block: the
comment on the load's guard keeps #3136's note on what it rests on, and
this PR's own diff is otherwise unchanged. Pre-commit, mypy included,
passes against main, and `tests/unit/torch/utils` passes on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-10-30 onto #3136 1491b66 (its MXFP4 refusal and
review fixes); this PR's own diff is unchanged. Pre-commit, mypy
included, passes against main, and `tests/unit/torch/utils` passes on
CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- GPU run, 2026-10-30, at #2743's head 7192e05, which contains this
PR at a4e2e1d: `tests/gpu/torch/utils/test_expert_parallel_load.py`
on transformers 5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

On 4 GPUs every case runs, the four with expert parallelism over half
the GPUs included, as at the run before. 1 GPU skips everything, so it
is no evidence.
- GPU run, 2026-10-30, at #2743's head 7c83230, which contains this
PR's head 4abc058:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

The same result as the run before, with the same test file, so it speaks
for the code underneath. The file loads an unquantized checkpoint on
transformers 5.18, so it cannot show the MXFP4 refusal or the block-FP8
floor of 5.8 firing: #3136's config-refusal table covers the MXFP4
refusal, and #2743's 5.7.0 test the block-FP8 floor, on CPU.
- Rebased 2026-10-30 onto #3136 ad9559d (meenchen's review, round 3,
on the base rebuilt with #2779 47ca6d1); this PR's own diff is
unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.
- Rebased 2026-10-30 onto main 059588b with #3136 eb6c382, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
3b814a6, #3067 fbff3f3 and #2740 186efae. #3136 now also
refuses a rank not yet bound to its GPU. This PR's own diff is
unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.
- Rebased 2026-10-30 onto #3136 56353ef, whose new commits change
only tests and one comment; this PR's own diff is unchanged. Pre-commit,
mypy included, passes against main, and `tests/unit/torch/utils` passes
on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- GPU run, 2026-10-30, at #2743's head 7540a2e, which contains this
PR at 6415c94: `tests/gpu/torch/utils/test_expert_parallel_load.py`
on transformers 5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

The same result as the runs before. On 4 GPUs every case runs, the four
with expert parallelism over half the GPUs included; 1 GPU skips
everything, so it is no evidence. Since that head, #3136 changed only
tests and one comment, so the result speaks for this head's code.
- Rebased 2026-10-30 onto #3136 75bc2d6, which trims the docstring
and comments this PR edits. This PR's own diff is unchanged but for its
docstring paragraph on the sharded dense part, now one sentence on the
trimmed docstring. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.
- Rebased 2026-10-30 with #3136 ab06860 onto main d464d30, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
8f9d877, #3067 d452768 and #2740 9a87d15; this PR's own diff
is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- GPU runs, 2026-10-30, at an earlier head with the same own diff
(#2743's 6f5391d, which contains this PR at 72f1e49), on
transformers 5.18.0:
  - `tests/gpu/torch/utils/test_expert_parallel_load.py`:

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

On 4 GPUs every case runs, the four with expert parallelism over half
the GPUs included; 1 GPU skips everything, so it is no evidence.
- `tests/unit/torch/utils/test_ep_loader.py` on a host with one visible
GPU: 59 of 59 passed.
- Review fixes, 2026-10-31 (cjluo-nv's review bot): the sharded dense
part materializes on the GPU each rank is bound to, not the one its
local rank numbers (f0c2a2b; the new test, local rank 1 bound to
cuda:0, fails without it), and the FSDP2 wrapping runs in the load's
failure-carrying guard, so a rank failing to wrap fails every rank
before DCP's collectives (0143941; the new bounded two-rank test
fails without it). Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.
- Rebased 2026-10-31 onto #3136 7747c5d, which moves
`test_ep_loader.py`'s cases into `test_model_load_utils.py` and the
`expert_parallel_mesh` ones into `test_moe_a2a.py`. This PR's test edits
moved with those cases, and its source diff is unchanged. Its
dense-sharding tests' experts class is now `_LayerExperts`, because
#2741's tests define an `_Experts` in the same file. Pre-commit, mypy
included, passes against main, and `tests/unit/torch/utils` collects and
passes the same cases as before on CPU, on transformers 5.18, 5.9.0 and
5.5.4.
- GPU run, 2026-10-31, at 0143941, which has this head's source:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

On 4 GPUs every case runs, the dense-sharded arms included. The file has
no dense-sharded arm with experts over half the GPUs, by design: this PR
refuses a dense part FSDP2-sharded within several expert-data-parallel
replicas (`edp > 1`), a layout the export cannot write yet.
- Review suggestion, 2026-10-31 (Claude): a transformers without
`get_total_byte_count` now warns that the load warmup keeps its own
sizing, as the other two fallbacks do (2860c0c); the test of that
release expects the warning and fails without it. Pre-commit, mypy
included, passes against main, and `tests/unit/torch/utils` passes on
CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- GPU check of the bound-device fix (f0c2a2b), 2026-10-31: 2 GPUs,
one visible per process, nccl, LOCAL_RANK 0/1, `DistributionPlan(2,
ep_size=2, dense="fsdp")`. The dense-sharded path ran on both ranks in
both runs.

  | Head | Rank 0 | Rank 1 |
  |---|---|---|
| 2860c0c (fixed) | loads | loads: its local rank is 1, and it
materializes on its bound GPU, cuda:0 |
| 591b986 (before the fix) | refuses, naming rank 1's failure |
fails: "CUDA error: invalid device ordinal" on cuda:1 |
- Rebased 2026-10-31 with #3136 930bb0f onto main e352cac, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
c78d2e5, #3067 adb9757 and #2740 dea9b79; this PR's own diff
is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- GPU run, 2026-10-31, at #2743's head 939bf6e on main e352cac,
which contains this PR's head 7ea4ac7:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

The same 12 cases as before, now over the rebuilt base, which includes
main's fused-MoE GPTQ change. On 4 GPUs every case runs, the
dense-sharded ones included: each rank's experts and the logits against
the whole model, and both quantized formats' training. 1 GPU skips
everything, so it is no evidence.
- Rebased 2026-11-01 with #3136 4dd38a3 onto main cf38ec3, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
c42cb61, #3067 cb83366 and #2740 ff372fb; this PR's own diff
is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-01 with #3136 35f1ac0 onto main 51645db, where
#3140 merged, on `shengliangx/hfep-base-6a-ep-loader` rebuilt there from
#2779 1aee07f, #3067 495f0be and #2740 6a06a54; this PR's
own diff is unchanged. The tree is identical to the previous head, as
#3140 merged as its own content. Pre-commit, mypy included, passes
against main.
- Rebased 2026-11-02 with #3136 39f55f2 onto main 2d4e8d0, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #3067
e6e422d and #2740 5829a85 (#2779 is now part of main); this PR's
own diff is unchanged. Pre-commit, mypy included, passes against main,
and `tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-02 with #3136 c0041d0 onto main 4f08a12, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2740
5829a85 (#3067 is now part of main); this PR's own diff is
unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-02 with #3136 d87f5bd onto main 653ef61, where
#2740 merged; this PR's own diff is unchanged. Pre-commit, mypy
included, passes against main, and `tests/unit/torch/utils` passes on
CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-03 onto main a6cc167, where #3136 merged; this
PR's own diff is unchanged. Pre-commit, mypy included, passes against
main.
- Review fix, 2026-11-03 (kevalmorabia97's review):
`test_a_tied_head_the_checkpoint_stores_alone_reaches_both_aliases`
writes copies of the tensors it read with `load_file` before `save_file`
overwrites their file (ab58a1a). With safetensors 0.8.0-rc.0 on
transformers 5.9.0 the test failed before the change, every loaded
tensor reading as zeros, and passes after it. Test only. Pre-commit,
mypy included, passes against main, and `tests/unit/torch/utils` and the
HF plugin tests pass on CPU at this head, on transformers 5.18, 5.9.0
and 5.5.4.
- Rebased 2026-11-03 onto main 7d44246, where #2780 merged, with the
junction the hf_ptq helper branch carried until now. In this PR's first
commit, the plan decision keeps #2780's refusal warning and drops its
give-way to whole-model FSDP2, for this PR's dense-part rule. A commit
on top (9768e7a) lifts `_route`'s refusal of an expert-parallel plan
with `dense="fsdp"`, says in `load_pretrained`'s docstring and the
distributed-loading guide that the dense part is replicated or
FSDP2-sharded as the plan says, updates #2780's plan tests to match (an
MoE left to decide keeps its EP degree and FSDP2-shards its dense part;
an expert-parallel plan with `dense="fsdp"` loads through
`load_and_prepare_ep`), and gives the dense-sharding tests' config
doubles a model type. The rest of this PR's diff is unchanged.
Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils`, `tests/unit/torch/export` and the HF plugin
tests pass on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4,
apart from 7 diffusers-export cases that fail on main itself in these
environments.
- Rebased 2026-11-03 onto main 5870698, where #3137 merged. In the
GPU test file, #3137's compressed-experts tests and this PR's
dense-sharded cases sit side by side. This PR's diff is otherwise
unchanged. Pre-commit, mypy included, passes against main, the GPU file
collects its 14 tests, and `tests/unit/torch/utils`,
`tests/unit/torch/export`, `tests/unit/torch/models` and the HF plugin
tests pass on CPU at this head, on transformers 5.18 and 5.5.4, apart
from 7 diffusers-export cases that fail on main itself in these
environments.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ Additive; the dense part stays
replicated unless asked.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A. The feature's entry lands with #2754.
- Did you get Claude approval on this PR?: Approved 2026-10-28 at
`05e21c2ba84`, and again on 2026-10-30 at `32d57f5c81c`, after the
review fixes in Testing, which change main's `_meta_dcp_load`.

### Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Expert-parallel model loading supports dense weights that are
replicated or FSDP-sharded, offering more flexible memory and
distribution options.
* FSDP2 wrapping follows model-defined layer boundaries, including
language-model layers in vision-language models.
* Block-FP8 checkpoints are refused by this PR's loader, whatever the
dense layout; #2743 adds their dequantization and keeps their dense part
replicated.
* **Bug Fixes**
* Improved loading and training consistency for expert-parallel models,
including quantized configurations.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…te the experts wrong (#2741)

> [!NOTE]
> **[3/3], split on review (2026-10-30):** this PR grew past the
500-line budget, so it is now three PRs, merged in order. #3131 [1/3]
holds the loader's helpers, #3136 [2/3] the expert-parallel loader, and
this PR, [3/3], the checks it runs once the experts load. #3131 and
#3136 have merged. The review history of the whole loader stays here.
>
> **Base:** `main`, where #3136 and #2780 have merged (2026-11-03). The
diff below is this PR's alone. #2780 merged before this PR after all, so
this PR now carries the test fixes the two need together.

<details><summary>Series layout</summary>

- `main` (#3131 [1/3], #3136 [2/3] and #2780 have merged)
  - **#2741 [3/3] the loader's post-load refusals (this PR)**
  - #2742 dense sharding → #2743 block-FP8 dequantization

</details>

### What does this PR do?

Type of change: New feature

Once the experts load, every rank checks what calibration relies on. It
refuses a load that would run without an error but compute or calibrate
wrong:
- **No match:** a plan whose expert containers match no module of the
loaded model, which would seed no expert state;
- **Not named as experts:** containers calibration would not take for
routed experts by name (`model_calib._is_routed_expert`). Their weight
scales would be merged over ranks holding different experts;
- **Count:** a count neither the config nor the experts' weights show
divisible by `ep_size`. transformers keeps the whole count on the
DTensors it shards the experts into, from 5.16;
- **Share:** a rank holding another share of the experts than `ep_size`
gives it, so the all-to-all transport would send tokens to ranks without
their experts. That happens to experts left whole on every rank, or
sharded over another mesh (a mesh-builder patch that stopped applying);
weights that show no share at all are refused too;
- **Not quantizable on the transport:** experts the grouped quantization
cannot run on the transport's forward, under any of three conditions:
  - they are stored transposed (GPT-OSS's), which it leaves unquantized;
- they have a quant wrapper of their own (Llama4's), whose forward drops
the transport's process group;
- their fused-experts wrapper quantizes on the eager forward only
(FP8's).

  ModelOpt's generic fused-experts wrapper is no bar.

The checks run inside #3136's `_refused_together` block around the load,
so a rank refusing alone fails every rank, and the share check needs no
collective of its own.

The transposed and eager-only rules are quantization's own. They become
one helper each in the HF quantization plugin (`_transposed_experts`,
`_eager_only_experts`), which `register_fused_experts_on_the_fly` and
`normalize_experts_impl_on_the_fly` now use too, so the load and
quantization cannot drift apart. The eager-only rule is the registry
lookup alone (#3067): the on-the-fly registration adds only the generic
wrapper.

### Usage

Automatic: `load_and_prepare_ep` (#3136) runs these checks after every
load.

### Testing

- `tests/unit/torch/utils/test_ep_loader.py` (CPU):
  - a plan matching no module;
  - containers without `experts` in their names;
- a count the config does not name, read off the weights, in these rows:
    - a DTensor count `ep_size` divides, and one it does not;
    - weights that show no count;
    - a rank sharded over another mesh;
    - plain weights left whole;
    - containers owning no weights.
- experts under a quant wrapper of their own, under the generic and the
FP8 fused-experts wrappers, and stored transposed by flag or by shape.
- **New:** `test_a_load_failing_on_one_rank_fails_every_rank` gains a
row where rank 1 alone holds half its share. Every rank raises rank 1's
refusal; with the load's guard made rank-local, the row fails.
- Mutation-checked when first written: each refusal fails without its
check.
- From before the split, on GPUs: #2742's quantized expert-parallel
training test, before this PR accepted the generic wrapper and with it:

| Dense layers, EP span, format | 2 GPUs: before, with the fix | 4 GPUs:
before, with the fix |
  |---|---|---|
  | replicated, all ranks, FP8 | pass, pass | pass, pass |
  | replicated, all ranks, NVFP4 | fail, pass | fail, pass |
  | replicated, half the ranks, FP8 | skipped | pass, pass |
  | replicated, half the ranks, NVFP4 | skipped | fail, pass |
  | FSDP2-sharded, all ranks, FP8 | pass, pass | pass, pass |
  | FSDP2-sharded, all ranks, NVFP4 | fail, pass | fail, pass |

- History of the checks before the split:
- 2026-10-27: the share check also reads plain (non-DTensor) weights,
and refuses weights that show no share.
- 2026-10-27: `_eager_only_experts` was factored out of the backend
normalization.
- 2026-10-28: the share check's comment says it reads the containers
together.
  - 2026-10-29: the eager-only rule took #3067's registry-only lookup.
- 2026-10-30, the split:
  - The loader and its helpers moved to #3136 and #3131.
- The checks moved inside #3136's load guard, and the share check's
`all_reduce` went away.
- `_eager_only_experts`, `_transposed_experts` and `_is_routed_expert`
are imported at module top. There is no import cycle: importing
`model_load_utils` already loads the quantization package, which never
imports it back. Checked in both orders.
  - Pre-commit, mypy included, passes against main.
- On CPU at this head, `tests/unit/torch/utils`, `test_fused_experts.py`
and the Nemotron-H tests pass on transformers 5.18 (880 passed) and on
5.5.4 (854 passed, 26 version-gated skips).
- Review fix, 2026-10-30 (Claude): the quantization plugin's rules and
`_is_routed_expert` are imported inside the functions that use them
again. Imported at the module top, a transformers the HF plugin cannot
load made `model_load_utils` unimportable, where the plugin's guarded
import only warns.
`test_the_loaders_import_without_the_hf_quantization_plugin` fails with
the imports at the top. Pre-commit, mypy included, passes against main.
The loader, fused-experts, Nemotron-H and `test_model_load_utils.py`
tests pass on CPU at this head, on transformers 5.18 and 5.5.4.
- Rebased 2026-10-30 onto #3136's review round, which now refuses before
the load a plan whose routed-expert container has no `grouped_gemm`
entry under it. The unnamed-container test's plan now shards its
`gate_up_proj` through `grouped_gemm`, as a real plan does, so it
reaches the refusal it checks. Pre-commit, mypy included, passes against
main. This PR's tests, with `test_model_load_utils.py`, pass on CPU on
transformers 5.18 and 5.5.4.
- Rebased 2026-10-30 with #3136 onto main 99c3834; this PR's own
diff is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18 and 5.5.4.
- Review fixes, 2026-10-30, round 2: the refusal of a plan matching no
module moved down to #3136, which needs it on its own. The
unquantizable-experts refusal takes its experts as a tuple, so a
one-shot iterable cannot skip its second rule;
`test_routed_experts_with_a_quant_wrapper_of_their_own_are_refused`
passes one and fails without the fix. The share check's comment says it
compares counts, and its test row is named for the mesh of another size
it builds. The two refusal tests' unplaced-keys stand-ins return a list,
as the function does. Pre-commit, mypy included, passes against main,
and `tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18 and 5.5.4.
- Rebased 2026-10-30 with #3136 onto the rebuilt
`shengliangx/hfep-base-6a-ep-loader`, which after #3131 merged holds
#2779 (on #3140, which moves `fsdp2_wrap` into
`modelopt/torch/utils/fsdp2_utils.py`), #3067 and #2740 on main
62a5c48. This PR's own diff is unchanged. Pre-commit, mypy included,
passes against main, and `tests/unit/torch/utils` passes on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-10-30 onto #3136 e251f6c (meenchen's review round).
The bounded two-rank test's docstring now names #3136's new `no-gpu` row
beside this PR's `share` row; this PR's own diff is otherwise unchanged.
Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Rebased 2026-10-30 onto #3136 1491b66 (its MXFP4 refusal and
review fixes); this PR's own diff is unchanged. Pre-commit, mypy
included, passes against main, and `tests/unit/torch/utils` passes on
CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-10-30 onto #3136 ad9559d (meenchen's review, round
3). #3136's tests now build their model with the shared
`model_with_decoder_at`, so this PR's quant-wrapper and
transposed-experts tests keep a small `_Experts` / `_SharedExperts` pair
of their own, which the quantization registry and the refusal messages
go by, and pass the experts directly rather than through a model. This
PR's own diff is otherwise unchanged. Pre-commit, mypy included, passes
against main, and `tests/unit/torch/utils` passes on CPU at this head,
on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-10-30 onto main 059588b with #3136 eb6c382, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
3b814a6, #3067 fbff3f3 and #2740 186efae. #3136 now also
refuses a rank not yet bound to its GPU. This PR's own diff is
unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.
- Rebased 2026-10-30 onto #3136 56353ef, whose new commits change
only tests and one comment; this PR's own diff is unchanged. Pre-commit,
mypy included, passes against main, and `tests/unit/torch/utils` passes
on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-10-30 with #3136 ab06860 onto main d464d30, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
8f9d877, #3067 d452768 and #2740 9a87d15; this PR's own diff
is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Rebased 2026-10-31 onto #3136 7747c5d, which moves
`test_ep_loader.py`'s cases into `test_model_load_utils.py` and the
`expert_parallel_mesh` ones into `test_moe_a2a.py`. This PR's test edits
moved with those cases, unchanged, and its source diff is unchanged.
Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` collects and passes the same cases as before on
CPU, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-10-31 with #3136 930bb0f onto main e352cac, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
c78d2e5, #3067 adb9757 and #2740 dea9b79; this PR's own diff
is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-01 with #3136 4dd38a3 onto main cf38ec3, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
c42cb61, #3067 cb83366 and #2740 ff372fb; this PR's own diff
is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-01 with #3136 35f1ac0 onto main 51645db, where
#3140 merged, on `shengliangx/hfep-base-6a-ep-loader` rebuilt there from
#2779 1aee07f, #3067 495f0be and #2740 6a06a54; this PR's
own diff is unchanged. The tree is identical to the previous head, as
#3140 merged as its own content. Pre-commit, mypy included, passes
against main.
- Rebased 2026-11-02 with #3136 39f55f2 onto main 2d4e8d0, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #3067
e6e422d and #2740 5829a85 (#2779 is now part of main); this PR's
own diff is unchanged. Pre-commit, mypy included, passes against main,
and `tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-02 with #3136 c0041d0 onto main 4f08a12, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2740
5829a85 (#3067 is now part of main); this PR's own diff is
unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-02 with #3136 d87f5bd onto main 653ef61, where
#2740 merged; this PR's own diff is unchanged. Pre-commit, mypy
included, passes against main, and `tests/unit/torch/utils` passes on
CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-03 onto main a6cc167, where #3136 merged; this
PR's own diff is unchanged. Pre-commit, mypy included, passes against
main.
- Review fix, 2026-11-03 (kevalmorabia97's review): the
`_transposed_experts` docstring says that a non-gated container
(`up_proj`, as Nemotron-H's) matches no `_fused_experts_wrapper`, so
only its `is_transposed` is read, and its own wrapper is responsible for
the layout (c0059c5). Docstring only. Pre-commit, mypy included,
passes against main.
- Rebased 2026-11-03 onto main 7d44246, where #2780 merged. This
PR's tests and #2780's in `test_model_load_utils.py` are kept side by
side: `_one_rank_fails_job`'s config doubles name `_NO_SPEC` with 8
experts, and the failure rows, cases and docstring of
`test_a_load_failing_on_one_rank_fails_every_rank` hold both PRs'
(66bcd7d). The stand-in configs of this PR's two new tests name
`_NO_SPEC` too (01b7fcf), as main's load now reads the model spec by
`model_type`; without it all 7 cases raised `AttributeError`. This PR's
source diff is unchanged. Pre-commit, mypy included, passes against
main, and `tests/unit/torch/utils`, `tests/unit/torch/export` and the HF
plugin tests pass on CPU at this head, on transformers 5.18, 5.9.0 and
5.5.4, apart from 7 diffusers-export cases that fail on main itself in
these environments.
- Rebased 2026-11-03 onto main 5870698, where #3137 merged. #3137's
tests and this PR's sit side by side in `test_model_load_utils.py`. A
test-only commit (df567f8) lets this PR's two new tests stand in for
`record_unplaced_source_keys` with keywords too, since the load now
passes `tensor_files=`; without it all 7 cases raised `TypeError`. This
PR's source diff is unchanged. Pre-commit, mypy included, passes against
main, and `tests/unit/torch/utils`, `tests/unit/torch/export`,
`tests/unit/torch/models` and the HF plugin tests pass on CPU at this
head, on transformers 5.18 and 5.5.4, apart from 7 diffusers-export
cases that fail on main itself in these environments.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ Additive; the HF plugin
refactor keeps quantization's behaviour.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A. The feature's entry lands with #2754.
- Did you get Claude approval on this PR?: Approved 2026-10-28 on the
unsplit PR; re-review requested 2026-10-30 for this slice.

### Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models. Merge order:
#3131 → #3136 → this PR. #2780, which routes `load_pretrained` to the
loader, has merged.

---------

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…rallel (#2743)

> [!NOTE]
> **Re-parented on 2026-10-30 with #2742 onto #3136**, #2741 [2/3] (the
expert-parallel loader). #2741 was split into three PRs on review: #3131
[1/3], #3136 [2/3] and #2741 [3/3], its post-load refusals. This PR
stays on #2742, whose dense-sharding rule it extends. The series first
became a tree on 2026-10-27 (this PR was [11/21]).
>
> **Base:** #2742, which is on `main`, where #3136 and #2780 have merged
(2026-11-03). The diff below is this PR's alone, and it merges after
#2742.
>
> **Re-parent:** this PR's diff is unchanged, except that its FP8
refusals run inside #3136's checks guard, so a rank refusing alone fails
every rank. That includes the per-tensor and unverifiable FP8-config
refusals, with their read of the checkpoint's `config.json`. The PR's 25
commits are condensed into one.

<details><summary>Series layout</summary>

- `main` (#3131 [1/3], #3136 [2/3] and #2780 have merged)
  - #2741 [3/3] the loader's post-load refusals
  - #2742 dense sharding
    - **#2743 block-FP8 dequantization (this PR)**

</details>

### What does this PR do?

Type of change: New feature

A native block-FP8 checkpoint, such as DeepSeek-V3, is dequantized to
bf16 while it loads expert-parallel, through transformers'
`FineGrainedFP8Config(dequantize=True)`. Quantizers attach to bf16
weights.

- **A block-FP8 checkpoint keeps its dense part replicated** (moved here
from #2742, where block FP8 did not load at all): DCP, which loads a
sharded dense part, cannot dequantize it: it reads the checkpoint as
stored. Left to decide, the plan (#2737) replicates it and says so; an
explicit `fsdp` is refused by the decision, naming
`MODELOPT_FSDP_SIZE=1`, and by the loader.
- **Per-tensor FP8 is refused.** AutoFP8's per-tensor checkpoints name
the same `quant_method` without a `weight_block_size`; transformers'
dequant would read them as 128x128 blocks and load their weights
unscaled, without an error. The load refuses them, before anything
loads, as `fp8-per-tensor`, naming the `weight_block_size` they lack. A
config handed in as transformers' FP8 config object, which fills in a
128x128 block, is read against the checkpoint's own `config.json`, and
refused when the path holds none to read (a Hub ID the load has not made
local yet). A tiny per-tensor FP8 Llama, loaded that way on transformers
5.18:

  | Weights | Largest absolute value |
  |---|---|
  | original | 0.079 |
  | loaded as block FP8 | 448 |

On transformers older than 5.10.1, the loader also patches that
dequantization in three ways:
- **Partial last block.** Stock dequantization infers the block size by
dividing the weight shape by the scale grid's, so it raises on a weight
whose last 128x128 block is partial, such as DeepSeek-V3's
`kv_a_proj_with_mqa`: (576, 7168) with a (5, 56) grid. It also raises on
an empty scale grid. The patch tiles each block's scale and crops to the
weight, using the checkpoint's `weight_block_size`.
- **Output dtype.** Stock dequantization does not take the destination
dtype. The patch emits the destination dtype the loader passes, else the
model's dtype.
- **Routed experts.** The FP8 quantizer prepends its dequantization to
every weight converter, so the fused-experts converter no longer starts
with a `MergeModulelist`. The loader then assigns no per-expert index,
and the expert-parallel shard slices each expert's rows instead of whole
experts: a DeepSeek-V3 `gate_up_proj` loads as (256, 32, 7168) and fails
calibration. The patch fuses the dequantization into the merge, so the
converter starts with a `MergeModulelist` again and runs as upstream
intends.

The patch installs once and survives a re-imported module. From
transformers 5.10.1 it keeps upstream's dequantization, which takes the
destination dtype and decodes E8M0 scales, and only pads a weight with a
partial last block to whole blocks for it, in its own fp8 dtype, since
upstream still infers the block from the scale grid: (448, 7168) under a
(4, 56) grid would otherwise be scaled in blocks of 112 rows.

The pre-5.10.1 dequantization scales its own fp32 copy of the weight in
place, through a view of its whole column blocks, against scales tiled
over the rows only, so no full-size fp32 tile is built. It is rebound
only once the expert converter's fuse is in, so a refused patch leaves
transformers as it was. Against scaling by a full-size fp32 tile, on CPU
with one thread:

| Weight | Peak memory growth: full tile | Peak memory growth: row-tiled
| Time: full tile | Time: row-tiled |
|---|---|---|---|---|
| (18432, 7168) | 1014 MiB (8.0 B/elem) | 759 MiB (6.0 B/elem) | 1.13 s
| 0.89 s |
| (7168, 2048) | 116 MiB | 88 MiB | 127 ms | 99 ms |

### Usage

Automatic: `load_and_prepare_ep` reads the checkpoint's
`quantization_config` and picks this path for `quant_method: fp8` with a
`weight_block_size`.

### Testing

- In `tests/unit/torch/utils/test_model_load_utils.py` (first written as
`test_fp8_dequant_patch.py`), on CPU:
- the patched dequant against a blockwise reference, a weight narrower
than one block included;
  - block size from the scale grid or the DeepSeek default;
  - the requested output dtype;
- an empty scale grid refused for a non-empty weight, rather than
passing it through unscaled;
  - idempotent installs that keep the latest dtype;
  - no double-wrapping of a re-imported module;
  - a clean undo;
- transformers 5.10.1 keeping its own dequant, padded for every
partial-block weight, a grid-divisible one included ((192, 256) on a (2,
2) grid);
- the fused expert converter, run on a real fused-experts converter: it
starts with a `MergeModulelist`, and per-expert block-FP8 weights come
out as the blockwise reference. It fails without the fusion, and if the
fused op skips the dequantization or the merge. Skipped on transformers
5.5, which has no `update_weight_conversions`.
- Mutation-checked: each test fails on at least one deliberate break of
the code.
- The block-size cases join the blockwise-reference table, and dtype and
version cases that took the same branch as a kept one are dropped.
- A refused fuse leaves transformers' dequant as it was.
- The block-FP8 dense layout: the loader refuses a plan that shards the
dense part (`test_model_load_utils.py`); the decision replicates it, or
refuses an explicit `fsdp`, naming `MODELOPT_FSDP_SIZE=1`
(`test_model_load_utils.py`).
- Per-tensor FP8 refused before anything loads and read as
`fp8-per-tensor` (`test_ep_loader.py`; both rows fail without the fix);
the block-FP8 configs in the tests carry their `weight_block_size`.
- Locally on this revision, the stack rebased on main c6fa613:
pre-commit over the whole PR against main, at every PR of the stack; and
at the stack's tip the full `tests/unit` on transformers 5.18 and 5.5
and on torch 2.8, and the hf_ptq example tests, each with the same
failures as main in that environment. CI runs the full unit suite.
- GPU results at a named head are under "GPU runs, 2026-10-30" at the
end of this section; the other GPU results here were measured on earlier
revisions of this branch.
- On GPUs: DeepSeek-V3 (block FP8) loads, calibrates and exports
expert-parallel through this path (see #2754).
- Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy
included, passes at this head over the whole change against main, and
the series' PRs merged together give a tree byte-identical to the former
stack's tip. This PR's own unit tests pass at this head on CPU (torch
2.14, transformers 5.18); CI runs the full unit suite.
- Review fixes, 2026-10-27: a dequant dtype the config names as a string
is read as the dtype it names; the test file's dead `Fp8Dequantize` skip
is gone. Each behaviour fix has a test that fails without it;
pre-commit, mypy included, and this PR's own unit tests pass at this
head.
- Review fixes, 2026-10-27, round 2: the block-FP8 case of the
replicated-dense-part test expects the out-of-memory warning #2742 now
gives when the plan replicates a dense part it cannot shard. Pre-commit,
mypy included, and this PR's own unit tests pass at this head.
- Review fixes, 2026-10-27, round 3: the pre-5.10.1 dequant override
takes the destination dtype the loader passes (5.10.1 is the first
release to pass one); a partial-block weight is padded in its own fp8
dtype; the 5.10.1 wrapper does not wrap itself on a re-import; a test
pins that transformers' config merge keeps the checkpoint's
`weight_block_size`; a comment says why the block-fp8 dense rule applies
under expert parallelism only; the per-tensor FP8 refusal names the
missing `weight_block_size`. Each behaviour fix has a test that fails
without it; pre-commit, mypy included, and this PR's own unit tests pass
at this head.
- Review fixes, 2026-10-27, round 4: a per-tensor FP8 checkpoint whose
config comes in as transformers' FP8 config object is refused too, its
block read from the checkpoint's own `config.json`; for a block-fp8
checkpoint, the several-replicas remedy names only
`MODELOPT_FSDP_SIZE=1`; both dequant patches check a weight's dims
before unpacking its shape, and the pre-5.10.1 override drops an FP4
branch it could not run. Each behaviour fix has a test that fails
without it; pre-commit, mypy included, and this PR's own unit tests pass
at this head.
- Review fixes, 2026-10-27, round 5: an FP8 config object whose
checkpoint has no readable `config.json` (a Hub ID not yet made local)
is refused rather than read as 128x128 blocks; for a block-fp8 dense
part within several replicas the remedy also names starting from an
unquantized copy. Each behaviour fix has a test that fails without it;
pre-commit, mypy included, and this PR's own unit tests pass at this
head.
- Rebased on 2026-10-27 onto main 7c3c868, after #2739 merged:
pre-commit, mypy included, and this PR's own unit tests pass at this
head, and the series' PRs merged together add no unit-test failure over
main (tests/unit/torch utils, export, quantization and models, on
transformers 5.18 and 5.5.4 and on torch 2.8).
- Re-cut on 2026-10-29 (see the note at the top): rebased onto the
re-cut #2742 with no source conflict. Its edits of the deleted
`test_resolve_distribution_plan.py` moved to `test_model_load_utils.py`.
Pre-commit, mypy included, passes at this head over the whole change
against main. On CPU at this head (torch 2.14), this PR's tests with
`test_ep_dense_sharding.py`, `test_ep_load_warmup.py`,
`test_ep_loader.py`, `test_meta_dcp_load.py` and
`test_model_load_utils.py` pass on transformers 5.18 and 5.5.4.
- Rebased 2026-10-29 onto #2742's review fixes (0350840) with no
conflict; this PR's own diff is unchanged. Pre-commit, mypy included,
passes against main, and `tests/unit/torch/utils` passes on CPU at this
head on transformers 5.18 and 5.5.4.
- Re-parented 2026-10-30 with #2742 onto #3136 (#2741 [2/3]). The FP8
refusals now run inside the loader's checks guard. Pre-commit, mypy
included, passes against main, and `tests/unit/torch/utils` passes on
CPU at this head, on transformers 5.18 and 5.5.4.
- Rebased 2026-10-30 with #2742 onto main 99c3834. The block-FP8
load notice now prints through `print_rank_0`, since #2742's loader no
longer imports `is_master`; the rest of this PR's diff is unchanged.
Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18 and 5.5.4.
- Rebased 2026-10-30 with #2742 onto #3136 96be72d; this PR's own
diff is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18 and 5.5.4.
- Review fix, 2026-10-30 (kevalmorabia97's review): below transformers
5.10.1, the block-FP8 override hands any weight that is not fp8 to
transformers' own dequant, in the model dtype. From 5.9 that dequant
unpacks packed FP4 (`float4_e2m1fn_x2` or `int8`, with `[1, 32]` scales,
as DeepSeek-V4 stores its routed experts), which the override returned
at half its width.
`test_a_weight_that_is_not_fp8_keeps_transformers_dequant` covers both
dtypes and fails without the fix. Pre-commit, mypy included, passes
against main, and `tests/unit/torch/utils` passes on CPU at this head,
on transformers 5.9.0, 5.18 and 5.5.4.
- Rebased 2026-10-30 onto #2742 201e2a5 (module-scope imports, a
hook-signature warning, one more remedy in a refusal); this PR's own
diff is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- GPU runs, 2026-10-30, at c44d5c1: NVIDIA GB200, torch 2.11.0a0
(nv26.02), transformers 5.9.0, pytest with `-p no:rerunfailures`.

  | Test file | 2 GPUs | 4 GPUs |
  |---|---|---|
| `tests/gpu/torch/utils/test_expert_parallel_load.py` | 8 passed, 4
skipped (123 s) | 12 passed (144 s) |

These results were taken at c44d5c1. Since then this PR moved its
unit tests into `test_model_load_utils.py` and was rebased with #2742,
last on 2026-10-30 onto the rebuilt base below; its source diff is
unchanged, and the file was re-run at this head (the last entry here).
- Rebased 2026-10-30 onto #2742 004de2a, which moved
`test_ep_dense_sharding.py` and `test_ep_load_warmup.py` into
`test_model_load_utils.py`; this PR's edit of the block-FP8 dense-layout
test moved with it, and its own diff is otherwise unchanged. Pre-commit,
mypy included, passes against main, and `tests/unit/torch/utils` passes
on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- 2026-10-30: `test_fp8_dequant_patch.py`, which covered
`model_load_utils` alone, moved into `test_model_load_utils.py` as one
section with its fixtures, marks and parametrization, as
CONTRIBUTING.md's naming rule asks (c884757). No test was added or
dropped: `tests/unit/torch/utils` gives the same counts on CPU on
transformers 5.18, 5.9.0 and 5.5.4, and pre-commit, mypy included,
passes against main.
- Rebased 2026-10-30 with #2742 onto the rebuilt
`shengliangx/hfep-base-6a-ep-loader`, which after #3131 merged holds
#2779 (on #3140, which moves `fsdp2_wrap` into
`modelopt/torch/utils/fsdp2_utils.py`), #3067 and #2740 on main
62a5c48. This PR's own diff is unchanged. Pre-commit, mypy included,
passes against main, and `tests/unit/torch/utils` passes on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4.
- GPU run, 2026-10-30, at this head, 9c917ae:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0, the test file unchanged since the run at c44d5c1.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

On 1 GPU every case skips, so it is no evidence. On 2 GPUs the four
cases with expert parallelism over half the GPUs skip; on 4 GPUs all 12
run: each rank's experts and the logits against the whole model, and FP8
and NVFP4 quantized training (scales and gradients), with the experts
over every GPU or half of them and the dense part replicated or
FSDP2-sharded.
- Rebased 2026-10-30 with #2742 onto #3136 e251f6c (meenchen's
review round); this PR's own diff is unchanged. Pre-commit, mypy
included, passes against main, and `tests/unit/torch/utils` passes on
CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- Review fix, 2026-10-30 (juhi10071998's comment): the expert-parallel
load refuses a block-fp8 checkpoint below transformers 5.8, inside its
checks guard. 5.7.0 has neither `Fp8Dequantize._dequantize_one` nor
`FineGrainedFP8HfQuantizer.update_weight_conversions`, so its experts
would load unscaled (7c83230).
`test_a_block_fp8_checkpoint_needs_the_transformers_that_dequantizes_its_experts`
presents 5.7.0 and fails without the check.
- Rebased 2026-10-30 with #2742 onto #3136 1491b66, which refuses
MXFP4: `_EP_SOURCE_FORMATS` is `{None, "fp8"}`, this PR's MXFP4 branch
went with #3136's, and the hub-id test lost its `mxfp4` row. Pre-commit,
mypy included, passes against main, and `tests/unit/torch/utils` passes
on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- GPU run, 2026-10-30, at 7192e05:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

On 4 GPUs every case runs, the four with expert parallelism over half
the GPUs included, as at the run before. 1 GPU skips everything, so it
is no evidence.
- GPU run, 2026-10-30, at this head, 7c83230:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

The same result as the run before, with the same test file, so it speaks
for the code underneath. The file loads an unquantized checkpoint on
transformers 5.18, so it cannot show the MXFP4 refusal or the block-FP8
floor of 5.8 firing:
`test_a_block_fp8_checkpoint_needs_the_transformers_that_dequantizes_its_experts`
(5.7.0) covers the block-FP8 floor, and #3136's `mxfp4` refusal row the
MXFP4 refusal, on CPU.
- Review suggestion, 2026-10-30 (Claude's review, approved): a
checkpoint `config.json` holding valid JSON other than an object reads
as one naming no FP8 block, so an FP8 config object meets the
`fp8-unverified` refusal instead of an `AttributeError` (5000023).
The new `config-json-not-an-object` row fails without the check.
Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4. The GPU run above was at 7c83230; this
one-line change is on a path the GPU test does not take.
- Rebased 2026-10-30 onto #2742 a8ba054, with #3136 ad9559d
(meenchen's review, round 3, on the base rebuilt with #2779
47ca6d1); this PR's own diff is unchanged. Pre-commit, mypy
included, passes against main, and `tests/unit/torch/utils` passes on
CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests
have not run on this head.
- Rebased 2026-10-30 onto #2742 6415c94, with #3136 eb6c382 on
main 059588b, on `shengliangx/hfep-base-6a-ep-loader` rebuilt there
from #2779 3b814a6, #3067 fbff3f3 and #2740 186efae. #3136
now also refuses a rank not yet bound to its GPU. This PR's own diff is
unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.
- Rebased 2026-10-30 onto #2742 02388f9, with #3136 56353ef,
whose new commits change only tests and one comment; this PR's own diff
is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- GPU run, 2026-10-30, at #2743's head 7540a2e, which contains this
PR's own commits as they are at this head:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

The same result as the runs before. On 4 GPUs every case runs, the four
with expert parallelism over half the GPUs included; 1 GPU skips
everything, so it is no evidence. Since that head, #3136 changed only
tests and one comment, so the result speaks for this head's code.
- Rebased 2026-10-30 onto #2742 72f1e49, with #3136 75bc2d6,
which trims the docstring and comments this PR edits. This PR's own diff
is unchanged but for its docstring lines, now on the trimmed docstring:
it loads block-FP8 checkpoints on 5.8 or later, and a sharded dense part
cannot load one. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.
- At b1cec70 on a host with one visible GPU,
`tests/unit/torch/utils/test_ep_loader.py` passed 57 of 57.
- Rebased 2026-10-30 onto #2742 591b986, with #3136 ab06860 on
main d464d30, on `shengliangx/hfep-base-6a-ep-loader` rebuilt there
from #2779 8f9d877, #3067 d452768 and #2740 9a87d15; this
PR's own diff is unchanged. Pre-commit, mypy included, passes against
main, and `tests/unit/torch/utils` passes on CPU at this head, on
transformers 5.18, 5.9.0 and 5.5.4.
- GPU runs, 2026-10-30, at an earlier head with the same own diff
(#2743's 6f5391d, which contains this PR's own commits as they are
at this head), on transformers 5.18.0:
  - `tests/gpu/torch/utils/test_expert_parallel_load.py`:

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

On 4 GPUs every case runs, the four with expert parallelism over half
the GPUs included; 1 GPU skips everything, so it is no evidence.
- `tests/unit/torch/utils/test_ep_loader.py` on a host with one visible
GPU: 59 of 59 passed.
- Rebased 2026-10-31 onto #2742 00a592d, with #3136 7747c5d,
which moves `test_ep_loader.py`'s cases into `test_model_load_utils.py`
and the `expert_parallel_mesh` ones into `test_moe_a2a.py`. This PR's
test edits moved with those cases, unchanged, and its own source diff is
unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` collects and passes the same cases as before on
CPU, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-10-31 onto #2742 7ea4ac7, with #3136 930bb0f on
main e352cac, on `shengliangx/hfep-base-6a-ep-loader` rebuilt there
from #2779 c78d2e5, #3067 adb9757 and #2740 dea9b79; this
PR's own diff is unchanged. Pre-commit, mypy included, passes against
main, and `tests/unit/torch/utils` passes on CPU at this head, on
transformers 5.18, 5.9.0 and 5.5.4.
- GPU run, 2026-10-31, at #2743's head 939bf6e on main e352cac,
which contains this PR's head:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

The same 12 cases as before, now over the rebuilt base, which includes
main's fused-MoE GPTQ change. On 4 GPUs every case runs, the
dense-sharded ones included: each rank's experts and the logits against
the whole model, and both quantized formats' training. 1 GPU skips
everything, so it is no evidence.
- Rebased 2026-11-01 onto #2742 dd04431, with #3136 4dd38a3 on
main cf38ec3, on `shengliangx/hfep-base-6a-ep-loader` rebuilt there
from #2779 c42cb61, #3067 cb83366 and #2740 ff372fb; this
PR's own diff is unchanged. Pre-commit, mypy included, passes against
main, and `tests/unit/torch/utils` passes on CPU at this head, on
transformers 5.18, 5.9.0 and 5.5.4.
- GPU run, 2026-11-01, at this head c278039 on main cf38ec3:
`tests/gpu/torch/utils/test_expert_parallel_load.py` on transformers
5.18.0.

  | GPUs | Passed | Skipped | Failed |
  |---|---|---|---|
  | 1 | 0 | 12 | 0 |
  | 2 | 8 | 4 | 0 |
  | 4 | 12 | 0 | 0 |

On 2 GPUs the four layouts with expert parallelism over half the GPUs
skip, by design. On 4 GPUs every case runs: each rank's experts and the
logits against the whole model, the backward pass, and FP8 and NVFP4
quantized training, across all three layouts. 1 GPU skips everything, so
it is no evidence.
- Rebased 2026-11-01 onto #2742 d3616e9, with #3136 35f1ac0 on
main 51645db, where #3140 merged, on
`shengliangx/hfep-base-6a-ep-loader` rebuilt there from #2779
1aee07f, #3067 495f0be and #2740 6a06a54; this PR's own diff
is unchanged. The tree is identical to the previous head, as #3140
merged as its own content. Pre-commit, mypy included, passes against
main.
- Rebased 2026-11-02 onto #2742 37d23c0, with #3136 39f55f2 on
main 2d4e8d0, on `shengliangx/hfep-base-6a-ep-loader` rebuilt there
from #3067 e6e422d and #2740 5829a85 (#2779 is now part of
main); this PR's own diff is unchanged. Pre-commit, mypy included,
passes against main, and `tests/unit/torch/utils` passes on CPU at this
head, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-02 onto #2742 91063d0, with #3136 c0041d0 on
main 4f08a12, on `shengliangx/hfep-base-6a-ep-loader` rebuilt there
from #2740 5829a85 (#3067 is now part of main); this PR's own diff
is unchanged. Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-02 onto #2742 a80bdb0, with #3136 d87f5bd on
main 653ef61, where #2740 merged; this PR's own diff is unchanged.
Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils` passes on CPU at this head, on transformers
5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-03 onto #2742 29e62bc, on main a6cc167, where
#3136 merged; this PR's own diff is unchanged. Pre-commit, mypy
included, passes against main.
- Review fix, 2026-11-03 (meenchen's review bot): below transformers
5.10.1, the install of a re-imported module replaces the dequant
override its earlier copy installed, around the method that override
replaced, rather than wrap it (0c67e21). Wrapped, a weight that is
not fp8 was cast to the earlier copy's dtype first, bfloat16 when unset,
and so rounded when the model loads at float32. The new re-import test
dequantizes a packed weight at float32 after a second install; it fails
without the change on transformers 5.18 and 5.9.0, and skips on 5.5.4,
which has no dequant of its own. The patch's `transformers` and
`packaging` imports and the load's `FineGrainedFP8Config` are
module-scope imports now. Pre-commit, mypy included, passes against
main, and the block-FP8 tests pass on CPU on transformers 5.18, 5.9.0
and 5.5.4.
- Rebased 2026-11-03 onto #2742 ab58a1a, with its test fix; this
PR's own diff is unchanged. Pre-commit, mypy included, passes against
main.
- Review fix, 2026-11-03 (Claude's review at 0c67e21): below
transformers 5.10.1 the model dtype the dequant override emits is kept
on `Fp8Dequantize` beside the override instead of in a module global
(ad10543), so a load through one copy of `model_load_utils` sets the
dtype the override another copy installed reads. The new test executes a
second copy of the module, installs the override through it, and loads
at float32 through the first copy; it gets bfloat16 without the change
on transformers 5.18 and 5.9.0. Pre-commit, mypy included, passes
against main, and `tests/unit/torch/utils` and the HF plugin tests pass
on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.
- Rebased 2026-11-03 onto #2742 9768e7a, on main 7d44246, where
#2780 merged. The load's refusals of an fp8 checkpoint it cannot take
now run after #2780's model-spec check, and this PR's test doubles name
`_NO_SPEC` where they conflicted with #2780's. A test-only commit on top
(7a1e112) gives the 5.8-floor test's double a model type, and makes
#2780's decision test expect the refusal of an fp8 checkpoint that names
no block size as `'fp8-per-tensor'`, since this PR loads block-fp8.
Pre-commit, mypy included, passes against main, and
`tests/unit/torch/utils`, `tests/unit/torch/export` and the HF plugin
tests pass on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4,
apart from 7 diffusers-export cases that fail on main itself in these
environments.
- Rebased 2026-11-03 onto #2742 c66e2ff, on main 5870698, where
#3137 merged. The load reads `_ckpt_quant_method(hf_config, ckpt_path)`
beside #3137's check for a checkpoint a model spec decompresses as it
loads. This PR's diff is otherwise unchanged. Pre-commit, mypy included,
passes against main, and `tests/unit/torch/utils`,
`tests/unit/torch/export`, `tests/unit/torch/models` and the HF plugin
tests pass on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4,
apart from 7 diffusers-export cases that fail on main itself in these
environments.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ Additive; runs only on the
expert-parallel load of a block-FP8 checkpoint.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A. The feature's entry lands with #2754.
- Did you get Claude approval on this PR?: ✅ Approved 2026-10-28 at
`a4609752e39`. The 2026-10-29 re-cut only moved its edits of a deleted
test file to `test_model_load_utils.py` (Testing).

### Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models; merges after
#2742.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added support for loading models from native block-FP8 checkpoints,
converting weights to the model’s configured dtype during loading.
* Improved handling of partial or empty scale grids, varied tensor
shapes, and per-expert weights.
* Added compatibility with older Transformers versions while preserving
built-in behavior in newer versions.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…3138)

> [!NOTE]
> **#2748 folded in (2026-11-03):** at the user's request to fold the
series backward, this PR now also carries #2748, the Kimi-K2.5 adapter
and `prepare` on the expert-parallel route only, with its commits as
they were. #2748 is closed and keeps its review history; its ticket is
this one, OMNIML-6078.
>
> **Split out of #2748** at the reviewers' request (2026-10-30), and
joined with it again: the Kimi-VL adapter, the patch it shares with the
Kimi-K2.5 adapter, the shared deepseek_v3 plan, and the Kimi-K2.5
adapter.
>
> **Base:** `main`, which holds what this PR builds on since they merged
on 2026-11-03: #3137 (on-the-fly decompression, with #2747 folded in)
and #2780 (the model specs, with #2746 folded in), and the remote-code
helpers (#3130). The branch `shengliangx/hfep-base-11b-kimi` that joined
#2780 and #3137 for this PR is retired.

<details><summary>Series layout</summary>

- `main` (#3130, #3131, #3132, #3134, #3135, #3140, #2779, #3074, #3067,
#2740, #3136, #2780, #3137, #2742, #2741 and #2743 have merged)
  - **#3138 Kimi-VL and Kimi-K2.5 adapters (this PR; #2748 folded in)**
    - #2749 [2/2] Kimi-K3 adapter

</details>

### What does this PR do?

Type of change: New feature

Kimi-VL ships its own modeling. It builds the language model's MoE from
an `nn.ModuleList` of per-expert Linears, which transformers'
expert-parallel machinery cannot shard. Its text backbone is
DeepSeek-V3, so the adapter swaps in transformers' native `deepseek_v3`
language model (fused `gate_up_proj`/`down_proj` experts, and the plan
to shard them). It keeps the bundled vision tower and projector from the
checkpoint's own code.

- **The patch Kimi-VL and Kimi-K2.5 share,** in
`modelopt/torch/models/_kimi.py`:
- **Refusals before anything changes:** transformers' own Kimi-K2.5 port
(from 5.13) of a quantized checkpoint, which the patch never reaches
(the port of an unquantized one loads as it is, below); a config with no
`text_config`; a missing checkpoint path; untrusted remote code
(`get_class_from_dynamic_module` runs the checkpoint's code without
checking `trust_remote_code`); and a quantization scheme on
`text_config` alone, which the native text config drops.
- **Import shims** for the symbols the checkpoint's modeling imports
that transformers 5 dropped (`is_torch_fx_available`, ...). They stub
only code paths the native language model never runs.
- **The swap itself:** the checkpoint's module-level
`DeepseekV3ForCausalLM` builds the native class from the native
`DeepseekV3Config` `prepare` puts on `text_config`, and the checkpoint's
own class from its own config. A `tie_weights` wrapper covers the
`recompute_mapping` newer transformers passes. A remote module that
defines no module-level `DeepseekV3ForCausalLM` stops the load instead
of passing for already patched.
- **Why the swap reads the config:** deciding the plan runs `prepare` on
a copy of the config on every rank before the plan is known, and the
patch lasts for the process. A plan that then keeps the load off expert
parallelism (replicated, or FSDP2) builds from the config as read, so it
gets the checkpoint's own language model, as with no patch. Before, the
native class was handed the checkpoint's own text config, which it
cannot take. The other effects of `prepare` are harmless off expert
parallelism: the import shims only add symbols the remote code imports,
the staging copies its sibling modules, the `tie_weights` wrapper only
drops keyword arguments for these untied checkpoints, and Kimi-K2.5's
decompressor rewrites only the configs it marked, which are the prepared
copies.
- **The swap is kept,** through `keep_patched` (#3130), on a module
transformers executes again or loads again under another name. The
reasons are in #3130's description.
- **Rank 0 only:** the adapters' notes and the shims' warnings print
there, labeled `kimi`.
- **The native backbone config:** `to_native_backbone_config`, in
`_onfly_decompress.py`, rebuilds the VLM's bundled text config as a
native transformers backbone config, carrying the decompression mark.
Both adapters' `prepare` call it; it moved here from #3137, which has no
caller of it.
- **The adapter is an `EPSpec.prepare` hook on the `kimi_vl` spec,** so
nothing outside `modelopt/torch/models` learns about Kimi.
- The spec registers only the expert-parallel section. The MoE and
export sections stay `deepseek_v3`'s, and sharing them is left to a
follow-up.
- The spec states the plan `prepare` attaches
(`EPSpec.prepared_ep_plan`, #2746): deepseek_v3's, now a shared
`BASE_MODEL_EP_PLAN`. So a run at an EP degree of 1, which skips
`prepare`, still finds the experts the checkpoint's own code builds, and
refuses to hold them whole on every rank.

<details><summary>The adapter's walkthrough (moved from
<code>kimi_vl/ep.py</code>'s module docstring)</summary>

```text
DP=EP a2a support for Kimi-VL-A3B (``kimi_vl``).

a remote-code VLM whose text backbone is a Moonlight (DeepSeek-V3-lite: MLA, ``q_lora_rank=None``, 64 routed experts, 2
shared) MoE and whose vision tower is MoonViT. The bf16 sibling of ``kimi_k25`` -- SAME shape, but NO compressed-tensors
decompressor (the checkpoint ships plain bf16 per-expert weights).

Challenge (identical to kimi_k25): the bundled ``modeling_kimi_vl.py`` builds its language model from
an INLINE ``DeepseekV3ForCausalLM`` whose ``DeepseekV3MoE`` is an ``nn.ModuleList`` of per-expert
``gate_proj``/``up_proj``/``down_proj`` Linears -- incompatible with transformers' EP machinery, which
needs the FUSED 3-D ``gate_up_proj``/``down_proj`` Parameters the transformers-native ``deepseek_v3``
modeling ships.

Solution (mirrors kimi_k25):
  1. Shim the import symbols ``modeling_kimi_vl`` needs that newer transformers dropped (e.g.
     ``is_torch_fx_available``) so the module IMPORTS -- they only guard torch.fx / legacy attn-mask
     paths that our native-LM + MoonViT forward never execute.
  2. Swap ``hf_config.text_config`` to a transformers-native ``DeepseekV3Config`` (drops the remote
     ``auto_map`` so the native class -- and its ``base_model_ep_plan`` -- is used for the LM).
  3. Apply the deepseek_v3 EP plan to that native ``text_config`` (post_init prefixes it to the VLM
     root as ``language_model.model.layers.*.mlp.experts.{gate_up_proj,down_proj}``).
  4. Monkey-patch ``modeling_kimi_vl.DeepseekV3ForCausalLM`` -> the transformers-native class, so
     ``KimiVLForConditionalGeneration.__init__`` (``self.language_model = DeepseekV3ForCausalLM(
     config.text_config)``) builds the EP-capable native LM while MoonViT + the KimiVL wrapper stay
     from the bundled code. Native deepseek_v3's per-expert WeightConverter fuses the checkpoint's
     per-expert ``.weight`` tensors into the 3-D params on load (no decompressor -- weights are bf16).

Version: no floor beyond the EP framework's 5.7; validated on transformers 5.9.0. The compat shims
stand in for the symbols the remote code imports that later releases dropped.
```

</details>

**The Kimi-K2.5 adapter** (folded in from #2748):

Kimi-K2.5 (and Kimi-K2.6 and Kimi-K2.7-Code, which share its model type)
is Kimi-VL's shape with compressed-tensors INT4 experts. Its adapter
swaps in transformers' native `deepseek_v3` language model as Kimi-VL's
does (above), through the patch they share, and decompresses the experts
as they load (#3137).

- **Install order.** The checkpoint's `config.json` keeps the
compressed-tensors config under `text_config`, and its config class also
copies it to the top level, which the expert-parallel load drops.
`prepare` installs the decompressor first of all, so a checkpoint the
decompressor refuses leaves the config and the loader untouched.
- **Which checkpoints decompress.** It installs the decompressor only
for a checkpoint whose `quantization_config` declares
compressed-tensors, as Kimi-K2.5, K2.6 and K2.7-Code all do.
  - A bf16 checkpoint of the model type loads as Kimi-VL's does.
- Any other scheme is left to the expert-parallel load's format
handling, which reads the copy of `quantization_config` the config class
keeps on the top level.
- A scheme on `text_config` alone, with no top-level copy, is refused
before anything changes.
- **transformers floor.** The spec declares the transformers 5.8 the
decompression needs.

**`prepare` now runs only on the expert-parallel route.** Before, the
plan decision ran a model spec's `prepare` on a copy of the config
before the plan was known, and the Kimi patches it made (the swapped
language-model class, the decompressor's loader hooks) stayed in the
process when the plan came out FSDP2-sharded, or whole and refused. Now:
- the decision reads the plan the spec states (`prepared_ep_plan`, on
`text_config` for both Kimi specs) and finds the experts through it in
the model the checkpoint's own code builds;
- before that it runs only `EPSpec.make_buildable`, what the
checkpoint's own code needs to import and build on the installed
transformers: the import shims, the staging of the code's two-deep
imports, and the VLM class's `tie_weights` adapted to the keyword
arguments `post_init` passes it. It changes nothing the model computes
or loads, warns rather than fails where it cannot adapt the class, and
patches no MoE or expert class, config, weight converter or
weight-loading hook;
- it asks the spec why `prepare` would refuse this checkpoint
(`EPSpec.prepare_refusal`), found without changing anything:
transformers' own Kimi-K2.5 port of a quantized checkpoint, code the
load may not run, experts the decompressor cannot take. With the EP
degree left to auto, such a checkpoint is planned without expert
parallelism, with the decision's warning;
- only an expert-parallel plan then runs `prepare`, on every rank, and
returns that config; any other plan returns the config as read.

`EPSpec.prepare`'s, `prepared_ep_plan`'s and
`resolve_distribution_plan`'s docs say so, replacing #2746's interim
contract. Kimi-VL's spec takes it too, and Kimi-K3 (#2749), on this PR,
uses it.

**transformers' own Kimi-K2.5 port.** transformers 5.13 added a
Kimi-K2.5 port, which a load without `trust_remote_code` builds. Its
language model is a native `deepseek_v3` with fused experts, and its
text config carries deepseek_v3's plan:
- **An unquantized checkpoint** read that way loads expert-parallel on
that plan without these patches, as it did before this PR: the adapter's
`prepare_refusal` and `prepare` leave its config as it is.
- **A quantized checkpoint** read that way is refused before anything
changes, naming `trust_remote_code=True` from a checkpoint whose
`auto_map` names its own `AutoConfig` as the way in: here only the
checkpoint's own code decompresses its experts.
- **Export names:** the port re-implements the vision tower and the
projector under other names. Its load-time conversions rename the vision
blocks' norms, MLPs and output projection and the projector, and split
the fused `wqkv` into q/k/v, so its export names differ from the
checkpoint's however it is distributed. Mapping them back is a question
for the export.

The spec records the 5.13 release.

**Why it patches the checkpoint's module, and again on every class
resolution:** see #3130, where `keep_patched` lives.
- The VLM class builds its language model from a name it looks up in its
module's globals.
- transformers re-executes or re-imports that module under the loaders
that resolve the class.

### Usage

Picked automatically by `apply_distributed_support` for `kimi_vl` and
`kimi_k25` configs under `torchrun`; see #2754.

### Testing

- `tests/unit/torch/models/kimi_vl/test_ep.py`, against a stand-in
remote module, covers:
- the text config rebuilt as a clean native `DeepseekV3Config` with the
EP plan;
- the remote modeling pointed at the native class, with the import shims
and the `tie_weights` wrapper;
- a load off expert parallelism, after `prepare` ran on a copy of its
config, building the checkpoint's own language model (it failed before
the fix);
- loud failures when the remote modeling did not load or has no
`DeepseekV3ForCausalLM`;
  - the refusal table, before anything is patched;
- the plan `prepare` attaches naming the same experts as the plan the
spec states.
- The shared tests are parametrized by adapter, Kimi-K2.5's cases
included.
- `tests/unit/torch/utils/test_model_load_utils.py`: a remote-code
Kimi-VL, as its own code builds it, is refused when left whole on every
rank at an EP degree of 1, with `prepare` never run.
- Review fix: the test is parametrized over the model type alone. Both
Kimi checkpoints build their experts at `.mlp.experts`.
- Locally at this head:
  - pre-commit against main passes, mypy included;
- `tests/unit/torch/models` and `test_model_load_utils.py` pass on CPU,
torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-10-30: rebased onto main 99c3834, through the new #2746,
#2747 (which now refuses compressed tensors the model leaves unplaced),
#3130 and #3132, with no conflicts and the own diff unchanged.
Pre-commit against main passes, mypy included;
`tests/unit/torch/models`, `test_model_load_utils.py` and
`test_ep_loader.py` pass on CPU, torch 2.14, on transformers 5.18 and
5.5.4.
- 2026-10-30, review round: docstrings only. The spec's says `is_moe`
misses the native `DeepseekV3MoE` blocks `prepare` installs until the
MoE and export sections are shared; the tie shim's says a tied
checkpoint still loads, with `lm_head.weight` reported missing (checked
on transformers 5.18 and 5.5.4). Pre-commit against main passes, mypy
included; `tests/unit/torch/models`, `test_model_load_utils.py` and
`test_ep_loader.py` pass on CPU, torch 2.14, on transformers 5.18 and
5.5.4.
- 2026-10-30: rebased onto the rebuilt base: main 7168107, where
#3130 and #3132 have merged, joined with #2746 5d51cd6 and #2747
4494480. There were no conflicts in this PR and its own diff is
unchanged. Pre-commit against main passes, mypy included. On CPU, torch
2.14, with transformers 5.18 and 5.5.4, these pass:
`tests/unit/torch/models`, `test_model_load_utils.py`,
`test_ep_loader.py`, `test_unplaced_source_keys.py`.
- 2026-10-30: the swap builds the native class only from the config
`prepare` made (review item from #2746). The new test fails without the
fix. Pre-commit against main passes. `tests/unit/torch/models`,
`test_model_load_utils.py` and `test_ep_loader.py` pass on CPU, torch
2.14, on transformers 5.18 and 5.5.4.
- 2026-10-30: rebased onto the rebuilt base: #2746 ffb58cf joined
with #2747 0f0570b, whose decompressor (#3137) now reverses its
converters as the ones it rewrote, so a decompressed checkpoint's export
keeps the checkpoint's names. No conflicts in this PR; the own diff is
unchanged. Pre-commit against main passes. `tests/unit/torch/models`,
`test_model_load_utils.py`, `test_ep_loader.py` and
`test_unplaced_source_keys.py` pass on CPU, torch 2.14: 572 on
transformers 5.18, and 447 passed, 56 skipped on 5.5.4.
- 2026-10-30: rebased onto the rebuilt base: #2746 d3a058e joined
with #2747 dc0a618, which walks a config's levels through one
`moe.config_levels` helper. #2746 now raises the generic left-whole
refusal as a ValueError and keeps the RuntimeError for a transformers
below the spec's floor; the remote-Kimi refusal test expects whichever
its transformers takes. Pre-commit against main passes, mypy included.
On CPU, torch 2.14, `tests/unit/torch/models` and
`tests/unit/torch/utils` pass on transformers 5.18 and 5.5.4.
- 2026-10-30: rebased onto the rebuilt base: main 059588b, with
#2746 0eadf4c joined with #2747 23eb10d. No conflicts in this
PR, and the own diff is unchanged. Pre-commit against main passes, mypy
included. On CPU, torch 2.14, `tests/unit/torch/models` and
`tests/unit/torch/utils` pass on transformers 5.18 and 5.5.4.
- 2026-10-31: rebased onto the rebuilt base: main d464d30, with
#2746 9267a0b joined with #2747 7d78ca1. #2746 added a decision
test where this PR adds the remote-Kimi stand-in and its refusal test;
both are kept, and the own diff is otherwise unchanged. #2748, on this
PR, now runs the Kimi-VL spec's `prepare` only for an expert-parallel
plan. Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils` pass on CPU,
torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-10-31, meenchen-nv's review:
- new 2-GPU case
`test_kimi_vl_loads_expert_parallel_and_matches_the_whole_model` writes
a tiny Kimi-VL-shaped remote checkpoint
(`tests/_test_utils/torch/kimi_vl_tiny_modeling.py`), with its experts
stored per expert, and loads it expert-parallel through the spec's
`prepare`. Each rank's native experts must equal the per-expert tensors
fused, the vision tower and projector must load as stored, nothing may
be left unplaced, and each rank's logits must match the whole model's on
one GPU. It passed on 2 GPUs and on 4 GPUs, with transformers 5.18.0, in
a run of the whole GPU file at #2749 183bf95, which contains this
PR. It loads through the spec's `prepare` and the loader directly, not
through the plan decision;
- the refusal of a scheme on `text_config` alone now names an
unquantized copy as the way in;
- the module docstring states the adapter's contract (the walkthrough is
above), `_kimi`'s module comment says why transformers is imported
inside its functions, and the tests move to
`tests/unit/torch/models/kimi_vl/test_ep.py`;
- rebased onto the rebuilt base: main d464d30, with #2746
77c4117 joined with #2747 7d78ca1. Pre-commit against main
passes, mypy included; on CPU, torch 2.14, `tests/unit/torch/models` and
`tests/unit/torch/utils` pass on transformers 5.18 and 5.5.4.
- 2026-10-31, Claude's review, at 3b5d7e2: the tiny Kimi-VL's
module-level `DeepseekV3ForCausalLM` is now the checkpoint's own,
per-expert class, so a swap that misses fails the GPU case's expert
checks; its CPU dry run passes. Pre-commit against main passes;
`tests/unit/torch/models` and `tests/unit/torch/utils` pass on CPU,
torch 2.14, on transformers 5.18 and 5.5.4. At #2749 4389363, which
contains this head, the whole GPU file passed at #2749 4389363, with
transformers 5.18.0 and compressed-tensors 0.19.0: on 2 GPUs, 5 passed
and `[ep-half]` skipped, as it needs 4; on 4 GPUs, all 6 passed. The
Kimi-VL case passed on both, now with the tiny model's own per-expert
language model, so a swap that missed would have failed it.
- 2026-10-31: rebased onto the rebuilt base: main e352cac, with
#2746 d8059b3 joined with #2747 6b87b69.
`test_model_load_utils.py` now holds the moved expert-parallel loader
tests; its one import conflict kept both sides' imports, and the own
diff is otherwise unchanged. At #2749 a60c01c, which contains this
head, the whole GPU file passed, with transformers 5.18.0 and
compressed-tensors 0.19.0: on 2 GPUs, 5 passed and `[ep-half]` skipped,
as it needs 4; on 4 GPUs, all 6 passed. The Kimi-VL case passed on both;
it loads through the spec and the loader directly, not through the plan
decision. Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils` pass on CPU,
torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-10-31, at cb90789: rebased onto the rebuilt base, whose #3137
now refuses a compressed-tensors older than 0.16; no conflicts, and the
own diff is unchanged. Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils` pass on CPU,
torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-10-31, at 7849516: rebased onto the rebuilt base, with #2747
dc7457e on #3137 81bcef7; no conflicts, and the own diff is
unchanged. Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils` pass on CPU,
torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-10-31: at #2749 76fcf09, which contains this PR at
7849516, the whole GPU file passed, with transformers 5.18.0,
compressed-tensors 0.20.1 and torch 2.15: on 2 GPUs, 5 passed and
`[ep-half]` skipped, as it needs 4; on 4 GPUs, all 6 passed. The Kimi-VL
case passed on both; it loads through the spec and the loader directly,
not through the plan decision.
- 2026-11-01, at e3d66fa: rebased onto the rebuilt base: main
cf38ec3, with #2746 997f9a6 joined with #2747 92830fc. #2746
now parametrizes `test_an_ep_capable_moe_may_be_fsdp2_sharded_instead`,
beside which this PR adds its Kimi tests; both are kept, and the own
diff is unchanged. Pre-commit against main passes, mypy included;
`test_model_load_utils.py`, `tests/unit/torch/models/kimi_vl` and
`test_ep_specs.py`, run in one process, pass on CPU, torch 2.14, on
transformers 5.18 and 5.5.4.
- 2026-11-01, at 4ff7f66: rebased onto the rebuilt base: main
51645db, with #2746 fce7da4 joined with #2747 1560551. The
base now carries #3135's check that a spec's `prepare` returns a config
with a plan where the expert-parallel load reads one; this PR's tests
pass with it. No conflicts; the own diff is unchanged. Pre-commit
against main passes, mypy included; `test_model_load_utils.py`,
`tests/unit/torch/models/kimi_vl` and `test_ep_specs.py`, run in one
process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-11-01: at #2749 0ea01fe, which contains this PR at
4ff7f66, the whole GPU file passed, with transformers 5.18.0 and
compressed-tensors 0.20.1: on 2 GPUs, 5 passed and `[ep-half]` skipped,
as it needs 4; on 4 GPUs, all 6 passed, in the file's order and again in
reverse, with the Kimi-K3 case first. The Kimi-VL case loads through the
spec and the loader directly, not through the plan decision.
- 2026-11-02, at 299afb6: rebased onto the rebuilt base: main
2d4e8d0, with #2746 55413f3 joined with #2747 9930028. In
that base the plan decision no longer refuses a model below the
transformers release its spec needs: it loads as a dense model does,
with a warning.
`test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error`
therefore skips below that release and expects only the ValueError above
it; it failed on transformers 5.5.4 before. The rest of the own diff is
unchanged. Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils`, run in one
process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-11-02, at 7eb48c0: rebased onto the rebuilt base: main
4f08a12, with #2746 e6efb94 joined with #2747 57daa15. No
conflicts; the own diff is unchanged. Pre-commit against main passes,
mypy included; `tests/unit/torch/models` and `tests/unit/torch/utils`,
run in one process, pass on CPU, torch 2.14, on transformers 5.18 and
5.5.4.
- 2026-11-02: at #2749 7cfcd54, which contains this PR at
7eb48c0, the whole GPU file passed, with transformers 5.18.0 and
compressed-tensors 0.20.1: on 2 GPUs, 5 passed and `[ep-half]` skipped,
as it needs 4; on 4 GPUs, all 6 passed, in the file's order and again in
reverse. The Kimi-VL case loads through the spec and the loader
directly, not through the plan decision.
- 2026-11-02, at 0b489fb: rebased onto the rebuilt base: main
653ef61, with #2746 f3503ff joined with #2747 3642bd0. No
conflicts; the own diff is unchanged. Pre-commit against main passes,
mypy included; `tests/unit/torch/models` and `tests/unit/torch/utils`,
run in one process, pass on CPU, torch 2.14, on transformers 5.18 and
5.5.4.
- 2026-11-03, at d24b5f5: #2748 folded in. This PR's branch was
fast-forwarded to #2748's head, so its commits are #2748's as they were,
on the same base. Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils`, run in one
process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4. On
GPU, the whole GPU file passed at #2749 7cfcd54, which contains this
code, with transformers 5.18.0 and compressed-tensors 0.20.1: on 2 GPUs,
5 passed and `[ep-half]` skipped, as it needs 4; on 4 GPUs, all 6
passed, in the file's order and again in reverse; its Kimi cases load
through the spec and the loader directly, not through the plan decision.
- 2026-11-03, at f5a0738, from Claude's review: the build step keeps
its `tie_weights` shim on a module transformers loads again
(235403a), and the plan decision asks the spec what its `prepare`
would refuse (`EPSpec.prepare_refusal`, f5a0738). So an EP degree
left to auto plans a checkpoint `prepare` refuses without expert
parallelism, with the warning naming why, rather than fail on every rank
in `prepare`. Each new test fails without its change. Pre-commit against
main passes, mypy included; `tests/unit/torch/models` and
`tests/unit/torch/utils`, run in one process, pass on CPU, torch 2.14,
on transformers 5.18 and 5.5.4.
- 2026-11-03, at 74450bb: rebased onto the rebuilt base, main
a6cc167 with #2780 c9f3a2d and #3137 25a4cd8.
`to_native_backbone_config` moved in from #3137, unchanged and with its
tests (e28e7e5). From Claude's review: the build step warns rather
than fails where it cannot adapt `tie_weights` (dc0c191), and its
test fails with the raise; the build step's comment names the module key
it shares with `prepare`'s swap (74450bb). Pre-commit against main
passes, mypy included; `tests/unit/torch/models` and
`tests/unit/torch/utils`, run in one process, pass on CPU, torch 2.14,
on transformers 5.18 and 5.5.4.
- 2026-11-03, at 3556f4d, from Claude's review:
`_onfly_decompress.decompression_refusal` states every check
`install_onfly_decompress` makes before it installs anything, now
including a config with nothing to decompress and a process already
decompressing a checkpoint compressed differently. Kimi-K2.5's
`prepare_refusal` asks it, and a test holds it to the install. The build
step's comment says why it runs for a decision that builds nothing
(b6bc19c). Pre-commit against main passes, mypy included;
`tests/unit/torch/models` and `tests/unit/torch/utils`, run in one
process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.
- 2026-11-03, at c6b6042, from Claude's review:
`decompression_refusal` reports the transformers 5.8 floor itself
(a6f48d3; its parity case fails without it), and one helper decides
whether Kimi-K2.5's `prepare` decompresses a checkpoint, for `prepare`
and `prepare_refusal` alike (c6b6042). Pre-commit against main
passes, mypy included; `tests/unit/torch/models` and
`tests/unit/torch/utils`, run in one process, pass on CPU, torch 2.14,
on transformers 5.18 and 5.5.4.
- 2026-11-03, at f6cdf9e: retargeted to #3137's branch and rebased
onto #3137 8feea6b (main 7d44246, which now holds #2780), off
the retired joining branch. No conflicts; the own diff is unchanged.
This PR now also runs on the #3137 commits the joining branch did not
carry: the off-index refusal, per-expert scales in the round trips, and
the tensor map read once and passed by keyword. Pre-commit against main
passes, mypy included; `tests/unit/torch/export`,
`tests/unit/torch/models` and `tests/unit/torch/utils`, run in one
process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4, apart
from seven `test_export_diffusers.py` cases that fail the same way on
main here, and `tests/unit/torch` passes without transformers.
- 2026-11-03, at 5d97c2c: rebased onto main 5870698, which now
holds #3137, keeping this PR's own commits alone; this PR now targets
main. From meenchen-nv's review, the decompressor's install and
`decompression_refusal` run one set of checks, which now includes the
raw dict (`_checked_compression`, da995de), and Kimi-K2.5's module
docstring is the adapter's contract, with both adapters' comments naming
the per-config build (caf4ea8). From Claude's review, the build step
warns on anything the checkpoint's code raises (2f461ac), and
`EPSpec.make_buildable` says that only a multi-rank decision runs it
(5d97c2c). Each behaviour change's test fails without it. Pre-commit
against main passes, mypy included; `tests/unit/torch/export`,
`tests/unit/torch/models` and `tests/unit/torch/utils`, run in one
process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4, apart
from seven `test_export_diffusers.py` cases that fail the same way on
main here, and `tests/unit/torch` passes without transformers.
- 2026-11-03, at 2c8a1ef: from meenchen-nv's review, reaffirmed,
transformers' own Kimi-K2.5 port of an unquantized checkpoint, whose
text config carries deepseek_v3's plan, is planned and loaded
expert-parallel on that plan, as before this PR, with `prepare` leaving
it as it is; the port of a quantized checkpoint is still refused
(5340e92). Its tests use the port as transformers 5.18 ships it and
fail without the change: the adapter's, and a decision that plans the
real port, built on the meta device, expert-parallel on a GPU the whole
model overfills but a rank's share fits. From CodeRabbit's review, the
unplaced-keys test imports transformers' loader modules at module scope
(2c8a1ef). Pre-commit against main passes, mypy included;
`tests/unit/torch/export`, `tests/unit/torch/models` and
`tests/unit/torch/utils`, run in one process, pass on CPU, torch 2.14,
on transformers 5.18 and 5.5.4, apart from seven
`test_export_diffusers.py` cases that fail the same way on main here,
and `tests/unit/torch` passes without transformers.
- 2026-11-03, at e2550e9: rebased onto main b002695, which now
holds #2742. The conflicts were context only: in
`test_model_load_utils.py` #2742 changed a neighbouring test's
parametrize, and in the GPU test it appended its own tests where this PR
appends the Kimi-VL one. The own diff is unchanged. Pre-commit against
main passes, mypy included; `tests/unit/torch/export`,
`tests/unit/torch/models` and `tests/unit/torch/utils`, run in one
process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4, apart
from seven `test_export_diffusers.py` cases that fail the same way on
main here, and `tests/unit/torch` passes without transformers.
- 2026-11-03, at 84b5029: from CodeRabbit's review, the Kimi-K2.5
port test says why it imports its config inside the test (f48c6eb).
From Claude's review, a CPU test pins that a load after the build step
holds the checkpoint as its own code does (84b5029): a tiny Kimi-VL
checkpoint, tied and untied, loads every stored tensor as stored, ties
its embedding exactly when the checkpoint does, and computes the same
logits as the checkpoint's own language model built with no patch; a
wrapper that skips the checkpoint's own `tie_weights` fails the tied
case. Pre-commit against main passes, mypy included; Kimi-VL's
`test_ep.py` passes on CPU, torch 2.14, on transformers 5.18 and 5.9,
and skips by its floor on 5.5.4 and without transformers.
- 2026-11-03, at 5e1bf0e: rebased onto main 3968c37, which now
holds #2741 and #2743. One conflict, an import line in
`test_model_load_utils.py`; the own diff is unchanged. Pre-commit
against main passes, mypy included; `tests/unit/torch/export`,
`tests/unit/torch/models` and `tests/unit/torch/utils`, run in one
process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4, apart
from seven `test_export_diffusers.py` cases that fail the same way on
main here, and `tests/unit/torch` passes without transformers.
- 2026-11-03, at #2749 d1f5293, which contains this PR's dc6a858
(a docstring since 5e1bf0e), the first GPU run under #2741's
post-load refusals: on a second GPU type, with transformers 5.18.0 and
compressed-tensors 0.20,
`tests/gpu/torch/utils/test_expert_parallel_load.py` passed, its Kimi-VL
and compressed-expert cases included, and no post-load refusal fired on
a Kimi checkpoint:

  | GPUs | Passed | Skipped |
  |---|---|---|
| 2 | 12 | 4: the half-width expert-parallel shapes, which need 4 GPUs |
  | 4 | 16 | 0 |
- Earlier history and GPU results: see #2748.

**Tests folded in from #2748:**

- `tests/unit/torch/models/kimi_k25/test_ep.py`, against a stand-in
remote module, covers:
- the text config rebuilt as a clean native `DeepseekV3Config` for
compressed-tensors, bf16 and block-FP8 experts, with the decompressor
installed for compressed-tensors ones only. The bf16 and block-FP8 cases
fail with an unconditional install.
- the decompressor installed first, while the scheme on `text_config` is
still present, and a decompressor refusal leaving `text_config` and the
remote-code staging untouched;
  - transformers 5.7 refused by the spec.
- The shared adapter tests
(`tests/unit/torch/models/kimi_vl/test_ep.py`) and the remote-Kimi
refusal test in `test_model_load_utils.py` gain their Kimi-K2.5 cases.
- The plan decision with `prepare` on the expert-parallel route only:
-
`test_only_an_expert_parallel_plan_patches_the_checkpoint_s_code[fsdp2|whole|expert-parallel]`
runs the real Kimi-VL and Kimi-K2.5 `prepare` through the decision; only
an expert-parallel plan patches the checkpoint's code;
- `test_a_checkpoint_its_spec_decompresses_is_decided_expert_parallel`:
a compressed-tensors Kimi-K2.5 checkpoint left to the decision is
planned expert-parallel, and a Kimi-VL one is not;
- the build step (`make_buildable`) lets the checkpoint's own VLM build
on transformers 5.18
(`test_the_build_step_lets_the_checkpoint_s_own_vlm_build_on_this_transformers`),
installs nothing when there is nothing to build, and runs at an EP
degree of 1 too
(`test_an_ep_degree_of_1_builds_the_checkpoint_s_own_vlm_to_decide`);
- `test_ep_specs.py`'s spec-field checks: a spec with a `prepare` states
its plan, and a build step, decompressed formats or a prepared plan each
come with a `prepare`.
- The decision's cost, measured on the real Kimi-K2.5 code and config
(61 layers, 384 routed experts) on the meta device, transformers 5.18,
one process; only rank 0 builds:

| The decision builds | Build time | Modules | Parameters | Peak host
RSS |
  |---|---|---|---|---|
| the prepared config: native DeepSeek-V3, fused experts (before) | 0.59
s | 1,455 | 1,250 | 0.62 GiB |
| the checkpoint's own code: per-expert experts (after) | 31.6 s |
116,655 | 70,310 | 5.16 GiB |

- Real-model runs, as #2748 recorded them (with the module-reload fix,
through the decompressor): Kimi-K2.7-Code, expert-parallel over 32 GPUs
(8 nodes, EP32), transformers 5.13.0, compressed-tensors 0.17.1, recipe
`nvfp4_experts_only_mse-kv_fp8_cast`, 4096 calibration samples of 4096
tokens, batch 1. Load, calibration and export completed:

  | Runs | Count | Result |
  |---|---|---|
| without the reload fix, first | 1 | failed before loading weights:
some ranks built the checkpoint's own `DeepseekV3ForCausalLM` on the
native config |
| without the reload fix, later | 2 | load, calibration and export
completed |
| with the reload fix | 2 | load, calibration and export completed,
through the decompressor |

- #2748's dated testing history stays on #2748.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ Additive. Registering the
Kimi-VL and Kimi-K2.5 specs reaches beyond the expert-parallel load only
in the plan decision of a multi-rank run: every rank runs the spec's
build step (`make_buildable`: the import shims, the staging of the
code's imports, and `tie_weights` adapted to the keyword arguments
`post_init` passes), which changes nothing the model computes or loads.
`prepare` runs only once the plan is decided expert-parallel. A run that
would leave an expert-parallel-capable MoE whole on every rank, with
nothing asking for that, is refused, as other expert-parallel MoEs are.
A single-process load is unchanged.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A. The feature's entry lands with #2754.
- Did you get Claude approval on this PR?: ❌ Not yet for the folded
change; requested on 2026-11-03.

### Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models: the Kimi-VL and
Kimi-K2.5 adapters, which #2748's split had made [1/2] (this PR) and
[2/2] (#2748), are one PR again.

Jira: OMNIML-6078, which #2748 shared.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
<!-- coderabbit-ai-generated:v1 -->

## Summary by CodeRabbit

* **New Features**
* Added expert-parallel loading support for Kimi-VL and Kimi-K2.5,
including compatible K2.6 and K2.7-Code checkpoints.
* Supported Kimi-K2.5 compressed checkpoints can now be decompressed
during expert-parallel loading.
* Kimi-VL vision components are retained while its language model uses
native DeepSeek-V3 support.
* **Improvements**
* Checkpoint compatibility is checked before expert-parallel
preparation. Automatic selection can fall back to standard loading with
a warning; explicitly requested expert parallelism reports why it cannot
proceed.
  * Preparation runs only when expert-parallel loading is selected.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>